Generación de paquetes de navegador de tamaño optimizado con TensorFlow.js

Visión general

TensorFlow.js 3.0 brinda soporte para crear paquetes de navegador orientados a la producción y de tamaño optimizado . Para decirlo de otra manera, queremos que le resulte más fácil enviar menos JavaScript al navegador.

Esta característica está dirigida a usuarios con casos de uso de producción que se beneficiarían particularmente de eliminar bytes de su carga útil (y, por lo tanto, están dispuestos a esforzarse para lograrlo). Para usar esta función, debe estar familiarizado con los módulos ES , las herramientas de agrupación de JavaScript, como el paquete web o el paquete acumulativo , y conceptos como la eliminación de códigos inactivos o sacudidas de árboles .

Este tutorial demuestra cómo crear un módulo tensorflow.js personalizado que se puede usar con un paquete para generar una compilación de tamaño optimizado para un programa que usa tensorflow.js.

Terminología

En el contexto de este documento, usaremos algunos términos clave:

Módulos ES : el sistema de módulos estándar de JavaScript . Introducido en ES6/ES2015. Identificable mediante el uso de declaraciones de importación y exportación .

Empaquetado : tomar un conjunto de activos de JavaScript y agruparlos/agruparlos en uno o más activos de JavaScript que se pueden usar en un navegador. Este es el paso que normalmente produce los activos finales que se entregan al navegador. Por lo general, las aplicaciones realizarán su propio empaquetado directamente desde las fuentes de la biblioteca transpiladas . Los paquetes comunes incluyen rollup y webpack . El resultado final de la agrupación se conoce como paquete (o, a veces, como un fragmento si se divide en varias partes).

Eliminación de Tree-Shaking/Dead Code : Eliminación de código que no se usa en la aplicación escrita final. Esto se hace durante la agrupación, normalmente en el paso de minificación.

Operaciones (Ops) : una operación matemática en uno o más tensores que produce uno o más tensores como salida. Las operaciones son código de "alto nivel" y pueden usar otras operaciones para definir su lógica.

Kernel : una implementación específica de una operación vinculada a capacidades de hardware específicas. Los kernels son de 'bajo nivel' y específicos de backend. Algunas operaciones tienen un mapeo uno a uno de la operación al núcleo, mientras que otras operaciones usan varios núcleos.

Alcance y casos de uso

Modelos gráficos de inferencia solamente

El caso de uso principal que escuchamos de los usuarios relacionados con esto, y que admitimos en esta versión, es el de hacer inferencias con los modelos gráficos de TensorFlow.js . Si está utilizando un modelo de capas de TensorFlow.js , puede convertirlo al formato de modelo gráfico mediante el convertidor tfjs . El formato del modelo gráfico es más eficiente para el caso de uso de inferencia.

Manipulación de tensor de bajo nivel con tfjs-core

El otro caso de uso que admitimos son los programas que usan directamente el paquete @tensorflow/tjfs-core para la manipulación de tensores de nivel inferior.

Nuestro enfoque para las construcciones personalizadas

Nuestros principios básicos al diseñar esta funcionalidad incluyen lo siguiente:

• Aproveche al máximo el sistema de módulos de JavaScript (ESM) y permita que los usuarios de TensorFlow.js hagan lo mismo.
• Haz que TensorFlow.js sea tan fácil de sacudir como sea posible por los paquetes existentes (por ejemplo, webpack, rollup, etc.). Esto permite a los usuarios aprovechar todas las capacidades de esos paquetes, incluidas funciones como la división de código.
• En la medida de lo posible, mantenga la facilidad de uso para los usuarios que no son tan sensibles al tamaño del paquete . Esto significa que las compilaciones de producción requerirán más esfuerzo, ya que muchos de los valores predeterminados en nuestras bibliotecas admiten la facilidad de uso en compilaciones optimizadas de tamaño.

El objetivo principal de nuestro flujo de trabajo es producir un módulo de JavaScript personalizado para TensorFlow.js que contenga solo la funcionalidad requerida para el programa que intentamos optimizar. Confiamos en los paquetes existentes para realizar la optimización real.

Si bien nos basamos principalmente en el sistema de módulos de JavaScript, también proporcionamos una herramienta CLI personalizada para manejar partes que no son fáciles de especificar a través del sistema de módulos en el código de cara al usuario. Dos ejemplos de esto son:

• Especificaciones del modelo almacenadas en archivos model.json
• El sistema de despacho de kernel específico de back-end que utilizamos.

Esto hace que generar una compilación de tfjs personalizada sea un poco más complicado que simplemente apuntar un paquete al paquete normal de @tensorflow/tfjs.

Cómo crear paquetes personalizados de tamaño optimizado

Paso 1: Determine qué núcleos está usando su programa

Este paso nos permite determinar todos los núcleos utilizados por los modelos que ejecuta o el código de procesamiento previo o posterior dado el backend que ha seleccionado.

Use tf.profile para ejecutar las partes de su aplicación que usan tensorflow.js y obtener los kernels. Se verá algo como esto

const profileInfo = await tf.profile(() => {
  // You must profile all uses of tf symbols.
  runAllMyTfjsCode();
});

const kernelNames = profileInfo.kernelNames
console.log(kernelNames);

Copie esa lista de núcleos en su portapapeles para el siguiente paso.

Debe perfilar el código usando los mismos backends que desea usar en su paquete personalizado.

Tendrá que repetir este paso si su modelo cambia o si cambia su código de procesamiento previo o posterior.

Paso 2. Escriba un archivo de configuración para el módulo tfjs personalizado

Aquí hay un archivo de configuración de ejemplo.

Se parece a esto:

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}

• kernels: La lista de kernels a incluir en el paquete. Copie esto de la salida del Paso 1.
• backends: la lista de backends que desea incluir. Las opciones válidas incluyen "cpu", "webgl" y "wasm".
• models: una lista de archivos model.json para los modelos que carga en su aplicación. Puede estar vacío si su programa no usa tfjs_converter para cargar un modelo gráfico.
• outputPath: una ruta a una carpeta para colocar los módulos generados.
• forwardModeOnly: configure esto en falso si desea incluir gradientes para los núcleos enumerados anteriormente.

Paso 3. Genere el módulo tfjs personalizado

Ejecute la herramienta de compilación personalizada con el archivo de configuración como argumento. Debe tener instalado el paquete @tensorflow/tfjs para tener acceso a esta herramienta.

npx tfjs-custom-module  --config custom_tfjs_config.json

Esto creará una carpeta en outputPath con algunos archivos nuevos.

Paso 4. Configure su paquete para alias tfjs en el nuevo módulo personalizado.

En paquetes como webpack y rollup podemos alias las referencias existentes a los módulos tfjs para apuntar a nuestros módulos tfjs personalizados recién generados. Hay tres módulos a los que es necesario crear un alias para obtener el máximo ahorro en el tamaño del paquete.

Aquí hay un fragmento de cómo se ve eso en el paquete web ( ejemplo completo aquí ):

...

config.resolve = {
  alias: {
    '@tensorflow/tfjs$':
        path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    '@tensorflow/tfjs-core$': path.resolve(
        __dirname, './custom_tfjs/custom_tfjs_core.js'),
    '@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
        __dirname, './custom_tfjs/custom_ops_for_converter.js'),
  }
}

...

Y aquí está el fragmento de código equivalente para el resumen ( ejemplo completo aquí ):

import alias from '@rollup/plugin-alias';

...

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
    },
  ],
}));

...

Si su paquete no admite la creación de alias de módulos, deberá cambiar sus instrucciones de import para importar tensorflow.js desde custom_tfjs.js generado que se creó en el Paso 3. Las definiciones de operaciones no se eliminarán, pero los núcleos seguirán siendo árboles. -agitado. En general, los núcleos que sacuden los árboles son los que proporcionan los mayores ahorros en el tamaño final del paquete.

Si solo usa el paquete @tensoflow/tfjs-core, solo necesita crear un alias para ese paquete.

Paso 5. Crea tu paquete

Ejecute su paquete (por ejemplo webpack o rollup ) para producir su paquete. El tamaño del paquete debe ser más pequeño que si ejecuta el paquete sin alias de módulo. También puede usar visualizadores como este para ver qué se convirtió en su paquete final.

Paso 6. Prueba tu aplicación

¡Asegúrate de probar que tu aplicación funciona como se espera!