Descripción general
Esta guía describe los mecanismos para definir operaciones personalizadas (ops), kernels y gradientes en TensorFlow.js. Su objetivo es proporcionar una descripción general de los conceptos principales y sugerencias de código que demuestran los conceptos en acción.
¿Para quién es esta guía?
Esta es una guía bastante avanzada que aborda algunos aspectos internos de TensorFlow.js y puede resultar particularmente útil para los siguientes grupos de personas:
- Usuarios avanzados de TensorFlow.js interesados en personalizar el comportamiento de diversas operaciones matemáticas (por ejemplo, investigadores que anulan implementaciones de gradientes existentes o usuarios que necesitan parchear funciones faltantes en la biblioteca)
- Usuarios que crean bibliotecas que amplían TensorFlow.js (por ejemplo, una biblioteca de álgebra lineal general construida sobre las primitivas de TensorFlow.js o un nuevo backend de TensorFlow.js).
- Usuarios interesados en contribuir con nuevas operaciones a tensorflow.js que quieran obtener una descripción general de cómo funcionan estos mecanismos.
Esta no es una guía para el uso general de TensorFlow.js, ya que aborda los mecanismos de implementación internos. No es necesario comprender estos mecanismos para utilizar TensorFlow.js.
Debe sentirse cómodo (o estar dispuesto a intentar) leer el código fuente de TensorFlow.js para aprovechar al máximo esta guía.
Terminología
En esta guía resulta útil describir algunos términos clave desde el principio.
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 específicas de hardware/plataforma. 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 utilizan múltiples kernels.
Gradient /GradFunc : la definición de 'modo inverso' de una operación/núcleo que calcula la derivada de esa función con respecto a alguna entrada. Los gradientes son código de 'alto nivel' (no específico del backend) y pueden llamar a otras operaciones o núcleos.
Registro del kernel : un mapa de una tupla (nombre del kernel, nombre del backend) a una implementación del kernel.
Registro de gradiente : un mapa de un nombre de kernel a una implementación de gradiente .
Organización del código
Las operaciones y los gradientes se definen en tfjs-core .
Los kernels son específicos del backend y se definen en sus respectivas carpetas del backend (por ejemplo, tfjs-backend-cpu ).
No es necesario definir operaciones, kernels y gradientes personalizados dentro de estos paquetes. Pero a menudo utilizará símbolos similares en su implementación.
Implementación de operaciones personalizadas
Una forma de pensar en una operación personalizada es simplemente como una función de JavaScript que devuelve alguna salida tensorial, a menudo con tensores como entrada.
- Algunas operaciones se pueden definir completamente en términos de operaciones existentes y deberían simplemente importar y llamar estas funciones directamente. Aquí hay un ejemplo .
- La implementación de una operación también se puede enviar a núcleos específicos de backend. Esto se hace a través de
Engine.runKernely se describirá con más detalle en la sección "implementación de núcleos personalizados". Aquí hay un ejemplo .
Implementación de kernels personalizados
Las implementaciones de kernel específicas del backend permiten una implementación optimizada de la lógica para una operación determinada. Los kernels son invocados por operaciones que llaman tf.engine().runKernel() . Las implementaciones de un kernel se definen por cuatro cosas.
- Un nombre de núcleo.
- El backend en el que está implementado el kernel.
- Entradas: argumentos tensoriales para la función del núcleo.
- Atributos: argumentos no tensoriales para la función del núcleo.
A continuación se muestra un ejemplo de una implementación del kernel . Las convenciones utilizadas para implementar son específicas del backend y se comprenden mejor al observar la implementación y la documentación de cada backend en particular.
Generalmente, los núcleos operan a un nivel inferior a los tensores y, en cambio, leen y escriben directamente en la memoria que eventualmente será envuelta en tensores por tfjs-core.
Una vez que se implementa un kernel, se puede registrar con TensorFlow.js utilizando la función registerKernel de tfjs-core. Puede registrar un kernel para cada backend en el que desee que funcione ese kernel. Una vez registrado, el kernel se puede invocar con tf.engine().runKernel(...) y TensorFlow.js se asegurará de enviarse a la implementación en el backend activo actual.
Implementación de gradientes personalizados
Los gradientes generalmente se definen para un kernel determinado (identificado por el mismo nombre de kernel usado en una llamada a tf.engine().runKernel(...) ). Esto permite a tfjs-core utilizar un registro para buscar definiciones de gradiente para cualquier kernel en tiempo de ejecución.
La implementación de gradientes personalizados es útil para:
- Agregar una definición de gradiente que puede no estar presente en la biblioteca
- Anulación de una definición de gradiente existente para personalizar el cálculo del gradiente para un núcleo determinado.
Puede ver ejemplos de implementaciones de gradientes aquí .
Una vez que haya implementado un gradiente para una llamada determinada, puede registrarlo con TensorFlow.js usando la función registerGradient de tfjs-core.
El otro enfoque para implementar gradientes personalizados que omite el registro de gradientes (y por lo tanto permite calcular gradientes para funciones arbitrarias de maneras arbitrarias es usar tf.customGrad .
A continuación se muestra un ejemplo de una operación dentro de la biblioteca que utiliza customGrad.