Generando paquetes de navegador de tamaño optimizado con TensorFlow.js

Descripción general

TensorFlow.js 3.0 brinda soporte para crear paquetes de navegadores 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 al reducir bytes de su carga útil (y, por lo tanto, están dispuestos a esforzarse para lograrlo). Para utilizar esta función, debe estar familiarizado con los módulos ES , las herramientas de agrupación de JavaScript, como webpack o rollup , y conceptos como la eliminación de códigos muertos 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 JavaScript estándar . Introducido en ES6/ES2015. Identificable mediante el uso de declaraciones de importación y exportación .

Agrupación : tomar un conjunto de activos de JavaScript y agruparlos/agregarlos en uno o más activos de JavaScript que se puedan utilizar en un navegador. Este es el paso que normalmente produce los recursos finales que se entregan al navegador. Las aplicaciones generalmente realizarán su propia agrupación directamente desde fuentes de biblioteca transpiladas . Los paquetes comunes incluyen rollup y webpack . El resultado final de la agrupación se conoce como paquete (o, a veces, como fragmento si se divide en varias partes).

Eliminación de código muerto o agitación de árboles : eliminación de código que no se utiliza 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 utilizar otras operaciones para definir su lógica.

Kernel : implementación específica de una opción vinculada a capacidades de hardware específicas. Los kernels son de "bajo nivel" y específicos del backend. Algunas operaciones tienen un mapeo uno a uno de la operación al kernel, mientras que otras usan múltiples kernels.

Alcance y casos de uso

Modelos gráficos de solo inferencia

El caso de uso principal que escuchamos de los usuarios relacionados con esto, y que respaldamos en esta versión, es el de realizar inferencias con modelos de gráficos de TensorFlow.js . Si está utilizando un modelo de capas de TensorFlow.js , puede convertirlo al formato de modelo de gráfico usando tfjs-converter . 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 JavaScript (ESM) y permita que los usuarios de TensorFlow.js hagan lo mismo.
• Haga que TensorFlow.js sea lo más sacudible posible mediante paquetes existentes (por ejemplo, paquete web, paquete acumulativo, 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 de nuestras bibliotecas admiten la facilidad de uso en comparación con las compilaciones de tamaño optimizado.

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

Si bien confiamos principalmente en el sistema de módulos 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 orientado al usuario. Dos ejemplos de esto son:

• Especificaciones del modelo almacenadas en archivos model.json
• La opción al sistema de distribución de kernel específico de backend que utilizamos.

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

Cómo crear paquetes personalizados de tamaño optimizado

Paso 1: determine qué núcleos utiliza su programa

Este paso nos permite determinar todos los kernels utilizados por cualquier modelo que ejecute o código de pre/postprocesamiento según el backend que haya seleccionado.

Utilice tf.profile para ejecutar las partes de su aplicación que utilizan tensorflow.js y obtener los núcleos. 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 a su portapapeles para el siguiente paso.

Debe crear un perfil del código utilizando los mismos backends que desea utilizar en su paquete personalizado.

Deberá repetir este paso si su modelo cambia o su código de pre/postprocesamiento cambia.

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 que se incluirán en el paquete. Copie esto del resultado del Paso 1.
• backends: la lista de backends que desea incluir. Las opciones válidas incluyen "cpu", "webgl" y "wasm".
• modelos: 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 de gráfico.
• OutputPath: una ruta a una carpeta para colocar los módulos generados.
• forwardModeOnly: configúrelo 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 poder 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 asignar el alias tfjs al nuevo módulo personalizado.

En paquetes como webpack y rollup, podemos asignar alias a las referencias existentes a los módulos tfjs para que apunten a nuestros módulos tfjs personalizados recién generados. Hay tres módulos a los que se les debe asignar un alias para lograr el máximo ahorro en el tamaño del paquete.

Aquí hay un fragmento de cómo se ve 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 el alias de módulo, deberá cambiar sus declaraciones 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 en forma de árbol, pero los núcleos seguirán siendo en árbol. -agitado. Generalmente, los granos que hacen temblar los árboles son los que proporcionan los mayores ahorros en el tamaño del paquete final.

Si solo está utilizando el paquete @tensoflow/tfjs-core, entonces solo necesita asignarle un alias a 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 menor que si ejecuta el paquete sin alias de módulo. También puedes usar visualizadores como este para ver qué se incluyó en tu paquete final.

Paso 6. Prueba tu aplicación

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