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.
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.
Deberás seguir los siguientes pasos importantes para implementar un modelo de TensorFlow Lite con entrenamiento integrado en el dispositivo:
Estos pasos se explican a continuación.
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:
# 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
@tf.function(input_signature=[tf.TensorSpec([None, IMG_SIZE, IMG_SIZE], tf.float32)]) def predict(self, x): return { "output": self.model(x) }
@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 }
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.
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:
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.
¡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.
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.