Lo nuevo en TensorFlow: Mejoras en KerasCV, DTensor, JAX2TF y Cuantización

Índice de contenido:

1. Introducción
2. Mejoras en KerasCV y KerasNLP
3. DTensor para modelos a gran escala
4. Integración de JAX en TensorFlow
5. Vista previa de la API de Cuantización de TensorFlow
6. Conclusiones
7. Preguntas frecuentes (FAQ)

[🌟 Sección 1: Introducción]

¡Hola y bienvenidos a TensorFlow y Keras en Google I/O! Soy Mani Varadarajan, el Director de Ingeniería de Google para las API de nuestro framework de aprendizaje automático. En esta sesión, vamos a echar un vistazo rápido a algunas de las muchas mejoras y adiciones que llegan este año a TensorFlow y la biblioteca de modelado de alto nivel Keras.

[🚀 Sección 2: Mejoras en KerasCV y KerasNLP]

Comencemos hablando de KerasCV y KerasNLP. Estas son bibliotecas modularizadas y potentes que te brindan acceso directo al estado del arte en visión por computadora y procesamiento de lenguaje natural. Ya sea que quieras clasificar imágenes, generar texto a partir de indicaciones o cualquier otra cosa en el medio, KerasCV y KerasNLP tienen todo lo que necesitas.

Estas bibliotecas proporcionan una base de vanguardia con solo unas pocas líneas de código. Y lo mejor de todo, al ser parte de Keras, están completamente integradas en el ecosistema de TensorFlow y respaldadas por una documentación de calidad. Entonces, puedes enfocarte en lo que más importa: innovar.

Aquí tienes un ejemplo rápido de código. Con menos de cinco líneas, puedes dar vida a lo último del estado del arte. Tenemos muchos más ejemplos como este, así que si quieres profundizar más, echa un vistazo a nuestra charla completa sobre KerasCV y KerasNLP, enlazada en la descripción a continuación.

[💪 Sección 3: DTensor para modelos a gran escala]

Ahora echemos un vistazo a DTensor y cómo te estamos permitiendo crear modelos de aprendizaje automático a una escala incomparable. Como todos hemos notado, los modelos se están volviendo enormes. A medida que los modelos crecen, entrenarlos y desplegarlos se vuelve aún más complejo.

Es aquí donde entra DTensor. Hoy, compartimos un conjunto de herramientas sencillas para construir modelos a una escala nunca antes vista, diseñadas para satisfacer las necesidades de todos en nuestra comunidad. DTensor está diseñado para ser flexible, eficiente y agnóstico al dispositivo.

Permíteme explicar cómo funciona. Tradicionalmente, los desarrolladores de aprendizaje automático han escalado sus modelos a través de la paralelización de datos, que implica dividir los datos y alimentarlos a instancias de modelo escaladas horizontalmente. Esto permite escalar el entrenamiento, pero con una gran limitación: todo debe caber dentro de un solo dispositivo.

A medida que los modelos se vuelven más grandes, esto ya no es una garantía. Sin embargo, con DTensor, puedes dividir de manera segura en múltiples dispositivos. Pero eso no es todo, DTensor también habilita la paralelización del modelo en sí, lo cual implica dividir el modelo y alimentar copias completas de los datos a cada fragmento.

Pero aquí está lo que hace a DTensor realmente especial. Con DTensor, puedes paralelizar tus datos y tu modelo en un solo lugar. Los principios detrás de la fragmentación de tu modelo para usar réplicas de datos o la fragmentación de tus datos para usar réplicas de modelo, siguen funcionando de la misma manera y se fusionan perfectamente bajo el poder de DTensor.

Todo lo que tienes que hacer es tomar cualquier modelo que ya hayas escrito y agregar unas pocas líneas para configurar adecuadamente y inicializar el contexto de DTensor. ¡Eso es todo! Si intentaras paralelizar datos de manera tradicional para el modelo en este ejemplo, probablemente daría un error. Los pesos del modelo son demasiado grandes para un solo dispositivo. Pero con DTensor, eso ya no es un problema. No tienes que preocuparte por reescribir tu modelo o escribir código diferente para diferentes estrategias de paralelización. Ya sea que estés utilizando un dispositivo o cien, DTensor te respalda. Y lo mejor es que sigue mejorando. El rendimiento actual ya está a la altura de los puntos de referencia de la industria de última generación, y ya tenemos mejoras planeadas para superar esos puntos de referencia. También queremos darte un adelanto rápido de lo que viene en el futuro. DTensor se integrará completamente con interfaces clave como tf.distribute y Keras en su conjunto, con un único punto de entrada independientemente del dispositivo y una serie de características que mejoran la calidad de vida. Si quieres aprender más, visita tensorflow.org o consulta las guías de integración en keras.io.

[🔬 Sección 4: Integración de JAX en TensorFlow]

Otra área que abordaremos hoy es lo que estamos haciendo para apoyar mejor a nuestra comunidad de investigación, especialmente con la aparición de JAX. Muchos de los avances en aprendizaje automático que ahora son bien conocidos tuvieron sus inicios en la investigación, incluidos los últimos modelos en las noticias como Bard y ChatGPT, que se basan en investigaciones publicadas por Google. JAX ha surgido como una herramienta confiable para gran parte de la investigación detrás de estos modelos. Sin embargo, llevar la investigación a producción es difícil. Es por eso que hemos estado reflexionando mucho sobre cómo podemos llevar el poder de producción del ecosistema de TensorFlow a JAX. Con ese fin, estamos muy emocionados de presentar JAX2TF, que proporciona un camino claro desde JAX hasta el ecosistema de TensorFlow.

JAX2TF es una API simple y liviana, y hay mucho que puedes hacer con ella. Por ejemplo, puedes llevar un modelo de JAX y desplegarlo en un servidor utilizando TF Serving o en un dispositivo utilizando TF Lite. También puedes usar JAX2TF para el ajuste fino, tomando un modelo preentrenado en JAX y continuar el entrenamiento desde allí. Incluso puedes fusionar modelos, tomando un modelo de JAX, combinándolo con capas adicionales y otros componentes, y luego entrenarlos como un solo modelo en TensorFlow. Esta es una capacidad muy poderosa, así que veamos cómo funciona. En términos generales, para llevar un modelo de JAX a TensorFlow, puedes hacer Algo como lo siguiente:

1. Primero, defines un modelo en JAX.
2. Segundo, creas una clase envolvente básica que utiliza jax2tf.convert para expresar los métodos de JAX como funciones de TensorFlow.
3. Y eso es prácticamente todo. Puedes guardarlo en un modelo guardado de TensorFlow y utilizarlo para servicio, ajuste fino, fusión, etc., tal como hemos explicado.

Es increíblemente fácil de hacer y tus modelos seguirán convergiendo rápidamente y con precisión. Usar JAX2TF es una forma poderosa de combinar JAX y TensorFlow para acelerar la investigación y llevarla a producción. Estamos emocionados de ver cómo la comunidad lo utiliza.

[⚡️ Sección 5: Vista previa de la API de Cuantización de TensorFlow]

Y para finalizar, permítenos mostrarte nuestras iniciativas para ayudarte a hacer que tus modelos sean más eficientes y fáciles de implementar en una amplia variedad de dispositivos. ¿Cómo lo estamos logrando? Con la API de Cuantización de TensorFlow, que estará disponible próximamente este año. La cuantización es un conjunto de técnicas que te permiten reducir el tamaño del modelo, lo que hace que se ejecuten más rápido y consuman menos recursos. Esto significa reducir la cantidad de memoria y cálculos necesarios para ejecutar los modelos. Esto puede reducir el consumo de batería en dispositivos móviles, así como la latencia del servidor y los costos de infraestructura.

Las versiones anteriores de las herramientas de cuantización, incluidas las de TensorFlow, eran limitadas. Sin embargo, la API de Cuantización de TensorFlow va mucho más allá de lo que se ha hecho antes. En primer lugar, es más flexible. Mientras que nuestra versión anterior estaba limitada a dispositivos móviles y requería el uso de TF Lite, esta API permite la cuantización en todas partes, incluyendo servidores, dispositivos móviles, dispositivos integrados, etc. En segundo lugar, es más fácil. Funciona de inmediato con configuraciones simples y no requiere cambios en el código del modelo. En tercer lugar, es más eficiente, brindándote el poder de cuantizar por capa, por operación o incluso por tensor, para construir de la manera que mejor te funcione.

Veámoslo en acción. Con solo unas pocas líneas de código, podemos preparar el modelo para la cuantización y entrenarlo o guardarlo dentro de un contexto de cuantización. En nuestras pruebas con el modelo MobileNet V2 en el Pixel 7, hemos logrado casi 17 veces más rendimiento en el rendimiento de servicio en comparación con la línea base de CPU no cuantizada. Pero aquí está lo mejor: todos estos beneficios vienen sin ningún impacto negativo notable en la precisión.

Y esto es solo el comienzo. La API de cuantización de TF aún no es pública, pero estará disponible pronto. Y seguiremos evolucionándola para brindar aún más beneficios. Y con esto, concluimos. Hoy te hemos mostrado solo algunas de las cosas clave en las que hemos estado trabajando, y hay mucho más por venir. Estamos emocionados de ver qué construirás, y siempre nos inspira la entusiasta participación y la asociación continua de nuestra comunidad. Gracias por tu atención.

FAQ (Preguntas frecuentes)

❓ ¿Qué es DTensor y qué ventajas ofrece?

DTensor es una herramienta que permite construir modelos de aprendizaje automático a gran escala, brindando flexibilidad, eficiencia y compatibilidad con diferentes dispositivos. Con DTensor, puedes escalar tus modelos tanto a través de paralelización de datos como de paralelización de modelo, todo en un solo lugar. Esto te permite entrenar y desplegar modelos masivos sin preocuparte por las limitaciones de un solo dispositivo.

❓ ¿Cuáles son las ventajas de utilizar la integración de JAX2TF?

JAX2TF ofrece una sencilla API para combinar los beneficios de JAX y TensorFlow. Con JAX2TF, puedes llevar modelos de JAX a TensorFlow y utilizarlos en servidores o dispositivos, sin necesidad de cambiar tu código de modelo. Esto acelera la investigación y facilita la transición de modelos de investigación a la producción.

❓ ¿Qué es la cuantización en TensorFlow y por qué es importante?

La cuantización en TensorFlow es un conjunto de técnicas que reducen el tamaño de los modelos de aprendizaje automático, lo que resulta en modelos más eficientes en términos de memoria y recursos computacionales. Esto permite que los modelos se ejecuten más rápido y consuman menos recursos, lo que es especialmente importante en dispositivos móviles y en entornos con recursos limitados. La API de Cuantización de TensorFlow amplía las capacidades de cuantización anteriores, ofreciendo más flexibilidad, facilidad de uso y eficiencia. Con la cuantización, puedes optimizar tus modelos sin comprometer su precisión.