Publicado por el equipo de TensorFlow Lite

TensorFlow Lite es el framework de aprendizaje automático de Google que se utiliza para implementar modelos de aprendizaje automático en diferentes dispositivos y superficies, como dispositivos móviles (iOS y Android), computadoras de escritorio y otros dispositivos periféricos. Hace poco, agregamos compatibilidad para que también se puedan ejecutar modelos de TensorFlow Lite en un navegador. Para crear apps con TensorFlow Lite, puedes utilizar un modelo existente de TensorFlow Hub o convertir un modelo de TensorFlow existente en un modelo de TensorFlow Lite mediante un conversor. Cuando se implementa un modelo en una app, puedes ejecutar un proceso de inferencia en el modelo según los datos ingresados.

Además de ejecutar un proceso de inferencia, TensorFlow Lite ahora admite el entrenamiento de modelos en un dispositivo. El entrenamiento en el dispositivo permite casos de uso de personalización interesantes donde se pueden ajustar los modelos según las necesidades del usuario. Por ejemplo, podrías implementar un modelo de clasificación de imágenes y permitir que un usuario lo optimice para reconocer especies de pájaros con el aprendizaje por transferencia, mientras le permites a otro usuario ajustar el mismo modelo de nuevo para reconocer frutas. Esta nueva función está disponible en la versión 2.7 y posteriores de TensorFlow y, en este momento, está disponible en las apps para Android. (En un futuro será compatible con iOS).

El entrenamiento en el dispositivo también es una base necesaria para los casos de uso de aprendizaje federado a fin de entrenar modelos globales sobre datos descentralizados. En esta entrada de blog, no hablamos sobre el aprendizaje federado, sino que nos centramos en ayudarte a integrar el entrenamiento en el dispositivo a tus apps para Android.

Luego, haremos referencia a una app de muestra de Colab y Android a medida que te guiamos en la ruta de implementación de extremo a extremo para el aprendizaje en el dispositivo a fin de ajustar un modelo de clasificación de imágenes.

Mejoras en cuanto al enfoque anterior

En nuestra entrada de blog de 2019, presentamos conceptos de entrenamiento en el dispositivo y un ejemplo en TensorFlow Lite. Sin embargo, hubo varias limitaciones. Por ejemplo, no era sencillo personalizar la estructura del modelo y los optimizadores. También existían varios modelos de TensorFlow Lite físicos (.tflite) en lugar de un solo modelo de TensorFlow Lite. Asimismo, no era sencillo almacenar y actualizar los pesos del entrenamiento. Como se explica a continuación, la última versión de TensorFlow Lite optimiza este proceso mediante opciones más convenientes para el entrenamiento en el dispositivo.

¿Cómo funciona?

Deberás seguir los siguientes pasos importantes para implementar un modelo de TensorFlow Lite con entrenamiento integrado en el dispositivo:

  • Crea un modelo de TensorFlow para entrenamiento e inferencia.
  • Convierte el modelo de TensorFlow al formato TensorFlow Lite.
  • Integra el modelo a la app para Android.
  • Invoca el entrenamiento del modelo en la app, de forma similar a como invocarías la inferencia del modelo.

Estos pasos se explican a continuación.

Crea un modelo de TensorFlow para entrenamiento e inferencia.

El modelo de TensorFlow Lite debería admitir tanto la inferencia como el entrenamiento del modelo, que por lo general implica guardar los pesos del modelo en el sistema de archivos y restaurarlos desde el mismo sistema de archivos. Se realiza esto para guardar los pesos del entrenamiento después de cada ciclo de entrenamiento, de modo que en el siguiente ciclo se puedan usar los pesos del anterior en lugar de comenzar un entrenamiento desde cero.

Sugerimos implementar estas tf.functions para que representen el entrenamiento, inferencia, pesos guardados y pesos cargados:

  • Una función de entrenamiento que entrena el modelo con datos de entrenamiento. La función de entrenamiento que se muestra a continuación realiza una predicción, calcula las pérdidas (o errores) y usa la clase tf.GradientTape() a fin de registrar operaciones para la diferenciación automática y actualizar los parámetros del modelo.
    # The `train` function takes a batch of input images and labels.
    @tf.function(input_signature=[
         tf.TensorSpec([None, IMG_SIZE, IMG_SIZE], tf.float32),
         tf.TensorSpec([None, 10], tf.float32),
     ])
    def train(self, x, y):
       with tf.GradientTape() as tape:
         prediction = self.model(x)
         loss = self._LOSS_FN(prediction, y)
       gradients = tape.gradient(loss, self.model.trainable_variables)
       self._OPTIM.apply_gradients(
           zip(gradients, self.model.trainable_variables))
       result = {"loss": loss}
       for grad in gradients:
         result[grad.name] = grad
       return result
    
  • Una función de inferencia o de predicción que invoca la inferencia del modelo. Esto es similar a cómo usas TensorFlow Lite para la inferencia.
    @tf.function(input_signature=[tf.TensorSpec([None, IMG_SIZE, IMG_SIZE], tf.float32)])
     def predict(self, x):
       return {
           "output": self.model(x)
       }
    
  • Una función de guardar/restablecer que guarda los pesos de entrenamiento (es decir, parámetros que utiliza el modelo) en formato de puntos de control en el sistema de archivos. El código de la función guardar se muestra a continuación.
    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
     def save(self, checkpoint_path):
       tensor_names = [weight.name for weight in self.model.weights]
       tensors_to_save = [weight.read_value() for weight in self.model.weights]
       tf.raw_ops.Save(
           filename=checkpoint_path, tensor_names=tensor_names,
           data=tensors_to_save, name='save')
       return {
           "checkpoint_path": checkpoint_path
       }
    

Convierte al formato TensorFlow Lite

Es posible que ya conozcas el flujo de trabajo para convertir el modelo de TensorFlow al formato TensorFlow Lite. Algunas de las funciones de bajo nivel para el entrenamiento en el dispositivo (p. ej., variables para almacenar los parámetros del modelo) aún son experimentales y otras (p. ej., la serialización del peso) en este momento se basan en operadores TF Select, por lo que deberás configurar estas marcas durante la conversión. Puedes descubrir ejemplos de todas las marcas que necesites establecer en Colab.

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL_DIR)
converter.target_spec.supported_ops = [
   tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
   tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.
]
converter.experimental_enable_resource_variables = True
tflite_model = converter.convert()

En Android, se puede realizar el entrenamiento en el dispositivo de TensorFlow Lite con las API de Java o C++. Puedes crear una instancia de la clase TensorFlow Lite Interpreter para cargar un modelo y así impulsar las tareas de entrenamiento del modelo. Antes, definimos varias tf.functions: se pueden invocar mediante la compatibilidad de TensorFlow Lite con las Firmas, lo que permite que un único modelo de TensorFlow Lite admita múltiples puntos de "entrada". Por ejemplo, definimos una función de entrenamiento para el entrenamiento en el dispositivo, que es una de las firmas del modelo. El entrenamiento se puede invocar mediante el método runSignature de TensorFlow Lite al especificar el nombre de la firma ("entrenamiento").

 /**
   * Runs one training step with the given bottleneck batches and labels.
   *
   * @param bottlenecks 2-D float array of bottleneck batches of size (BATCH_SIZE, BOTTLENECK_SIZE)
   * @param labels 2-D float array of label batches of size (BATCH_SIZE, NUM_CLASSES)
   * @return the training loss
   */
  float runTraining(float[][] bottlenecks, float[][] labels) {
    Map inputs = new HashMap'<>'();
    inputs.put("bottleneck", bottlenecks);
    inputs.put("label", labels);
 
    Map outputs = new HashMap'<>'();
    FloatBuffer loss = FloatBuffer.allocate(1);
    outputs.put("loss", loss);
 
    this.interpreter.runSignature(inputs, outputs, "train");
 
    return loss.get(0);
  }

De forma similar, el siguiente ejemplo muestra cómo invocar la inferencia mediante la firma "inferencia" del modelo.

/**
   * Invokes inference on the given image batches.
   *
   * @param testImage 3-D float array of image of size (IMG_SIZE, IMG_SIZE, 3)
   * @return 1-D float array of softmax output of prediction
   */
  float[] runInference(float[][][] testImage) {
    // Run the inference.
    Map inputs = new HashMap<>();
    inputs.put("feature", new float[][][][] {testImage});

    Map outputs = new HashMap<>();
    float[][] output = new float[1][numClasses];
    outputs.put("output", output);

    this.interpreter.runSignature(inputs, outputs, "infer");
    return output[0];
  }

¡Eso es todo! Ahora tienes un modelo de TensorFlow Lite que puede utilizar entrenamiento en el dispositivo. Esperamos que esta guía sobre el código te dé una buena idea sobre cómo ejecutar el entrenamiento en el dispositivo en TensorFlow Lite. Nos entusiasma ver hasta dónde lo llevas.

Consideraciones prácticas

En teoría, deberías poder aplicar el entrenamiento en el dispositivo en TensorFlow Lite a cualquier caso de uso que admita TensorFlow. Sin embargo, en realidad existen muy pocas consideraciones prácticas que necesitas tener en cuenta antes de implementar el entrenamiento en el dispositivo en tus apps y son las siguientes:

  • Casos de uso: El ejemplo de Colab muestra un ejemplo de entrenamiento en el dispositivo para un caso de uso sobre la visión. Si tienes problemas en modelos o casos de uso específicos, cuéntanos en GitHub.
  • Rendimiento: Según el caso de uso, el entrenamiento en el dispositivo podría tardar desde unos segundos hasta mucho más tiempo. Si ejecutas el entrenamiento en el dispositivo como parte de una función del usuario (p. ej., tu usuario final interactúa con la función), debes medir el tiempo necesario para una gran variedad de posibles entradas de entrenamiento en tu app a fin de limitar el tiempo de entrenamiento. Si tu caso de uso requiere tiempos muy largos de entrenamiento en el dispositivo, primero considera entrenar un modelo mediante una computadora de escritorio o la nube y luego ajústalo en el dispositivo.
  • Uso de la batería: Al igual que la inferencia de modelos, invocar un entrenamiento de modelos en el dispositivo puede provocar que se agote la batería. Si el entrenamiento de modelos forma parte de una función que no es para el usuario, recomendamos seguir las pautas de Android para implementar tareas en segundo plano.
  • Entrenar desde cero o volver a entrenar: En teoría,debería ser posible entrenar un modelo desde cero en el dispositivo con las funciones mencionadas anteriormente. Sin embargo, en realidad,el entrenamiento desde cero involucra una gran cantidad de datos de entrenamiento y puede tardar varios días, incluso en servidores con procesadores potentes. Como consecuencia, para las aplicaciones en el dispositivo recomendamos volver a realizar el entrenamiento en base a un modelo ya entrenado (es decir, realizar una transferencia de aprendizaje) como se mostró en el ejemplo de Colab.

Guía

Los trabajos futuros incluyen (pero no se limitan a) compatibilidad de entrenamiento en el dispositivo en iOS, mejoras de rendimiento a fin de aprovechar los aceleradores en el dispositivo (p. ej., GPU) para el entrenamiento en el dispositivo, lo que reduce el tamaño binario al implementar más operaciones de entrenamiento de forma nativa en TensorFlow Lite, compatibilidad de API de nivel superior (p. ej., a través de la biblioteca de tareas de TensorFlow Lite) para abstraer los detalles de implementación y ejemplos que cubren otros casos de uso de entrenamiento en el dispositivo (p. ej., PLN). En nuestra guía a largo plazo, se intentan brindar soluciones de aprendizaje federado de extremo a extremo en el dispositivo.

Próximos pasos

¡Gracias por leernos! Nos entusiasma ver lo que creas con el aprendizaje en el dispositivo. Una vez más, aquí están los enlaces a la app de muestra y de Colab. Si tienes algún comentario, compártelo con nosotros en el Foro de TensorFlow o en GitHub.

Agradecimientos

En esta publicación, se reflejan las contribuciones significativas de muchas personas en el equipo de TensorFlow Lite de Google, incluidos Michelle Carney, Lawrence Chan, Jaesung Chung, Jared Duke, Terry Heo, Jared Lim, Yu-Cheng Ling, Thai Nguyen, Karim Nosseir, Arun Venkatesan, Haoliang Zhang, otros miembros del equipo de TensorFlow Lite y nuestros colaboradores en Google Research.