Escribir operaciones, kernels y degradados personalizados en TensorFlow.js,Escribir operaciones, kernels y degradados personalizados en TensorFlow.js

Visió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 principales conceptos y sugerencias para el código que demuestran los conceptos en acción.

¿Para quién es esta guía?

Esta es una guía bastante avanzada que toca algunas partes internas de TensorFlow.js, puede ser particularmente útil para los siguientes grupos de personas:

  • Usuarios avanzados de TensorFlow.js interesados ​​en personalizar el comportamiento de varias operaciones matemáticas (por ejemplo, investigadores que anulan las implementaciones de gradientes existentes o usuarios que necesitan parchear la funcionalidad que falta en la biblioteca)
  • Usuarios que crean bibliotecas que amplían TensorFlow.js (p. ej., 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 desean 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 se incluye en los mecanismos de implementación internos. No necesita comprender estos mecanismos para usar TensorFlow.js

Debe sentirse cómodo (o dispuesto a intentarlo) leyendo el código fuente de TensorFlow.js para aprovechar al máximo esta guía.

Terminología

Para esta guía, es útil describir algunos términos clave por adelantado.

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 específicas de hardware/plataforma. 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.

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 degradados son código de "alto nivel" (no específicos del backend) y pueden llamar a otras operaciones o kernels.

Registro del kernel : un mapa de una tupla (nombre del kernel, nombre del backend) a una implementación del kernel.

Registro de degradado: un mapa de un nombre de kernel a una implementación de degradado .

Organización del código

Las operaciones y los gradientes se definen en tfjs-core .

Los núcleos son específicos de backend y se definen en sus respectivas carpetas de backend (p. ej ., tfjs-backend-cpu ).

No es necesario definir operaciones, kernels y gradientes personalizados dentro de estos paquetes. Pero a menudo usará 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 una salida de tensor, a menudo con tensores como entrada.

  • Algunas operaciones se pueden definir completamente en términos de operaciones existentes, y solo deben importar y llamar a estas funciones directamente. Aquí hay un ejemplo .
  • La implementación de una operación también puede enviarse a kernels específicos de backend. Esto se hace a través Engine.runKernel y se describirá con más detalle en la sección "implementación de kernels personalizados". Aquí hay un ejemplo .

Implementación de núcleos personalizados

Las implementaciones de kernel específicas de back-end permiten una implementación optimizada de la lógica para una operación determinada. Los núcleos son invocados por operaciones que llaman a tf.engine().runKernel() . Las implementaciones de un núcleo se definen por cuatro cosas

  • Un nombre de núcleo.
  • El backend en el que se implementa el kernel.
  • Entradas: Argumentos de tensor para la función kernel.
  • Atributos: Argumentos no tensoriales para la función kernel.

Aquí hay un ejemplo de una implementación del núcleo . Las convenciones que se usan para implementar son específicas del backend y se entienden mejor al observar la implementación y la documentación de cada backend en particular.

Generalmente, los núcleos operan a un nivel más bajo que los tensores y, en su lugar, leen y escriben directamente en la memoria que tfjs-core eventualmente envolverá en tensores.

Una vez que se implementa un kernel, se puede registrar con TensorFlow.js usando 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 enviar a la implementación en el backend activo actual.

Implementación de degradados personalizados

Los gradientes generalmente se definen para un kernel determinado (identificado por el mismo nombre de kernel utilizado en una llamada a tf.engine().runKernel(...) ). Esto permite que tfjs-core use un registro para buscar definiciones de gradiente para cualquier kernel en tiempo de ejecución.

La implementación de degradados personalizados es útil para:

  • Agregar una definición de gradiente que puede no estar presente en la biblioteca
  • Anular una definición de gradiente existente para personalizar el cálculo de gradiente para un kernel determinado.

Puede ver ejemplos de implementaciones de degradado aquí .

Una vez que haya implementado un gradiente para una llamada determinada, puede registrarse con TensorFlow.js usando la función registerGradient de tfjs-core.

El otro enfoque para implementar gradientes personalizados que pasa por alto el registro de gradientes (y, por lo tanto, permite calcular gradientes para funciones arbitrarias de formas arbitrarias es usar tf.customGrad .

Aquí hay un ejemplo de una operación dentro de la biblioteca de usar customGrad

,

Visió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 principales conceptos y sugerencias para el código que demuestran los conceptos en acción.

¿Para quién es esta guía?

Esta es una guía bastante avanzada que toca algunas partes internas de TensorFlow.js, puede ser particularmente útil para los siguientes grupos de personas:

  • Usuarios avanzados de TensorFlow.js interesados ​​en personalizar el comportamiento de varias operaciones matemáticas (por ejemplo, investigadores que anulan las implementaciones de gradientes existentes o usuarios que necesitan parchear la funcionalidad que falta en la biblioteca)
  • Usuarios que crean bibliotecas que amplían TensorFlow.js (p. ej., 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 desean 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 se incluye en los mecanismos de implementación internos. No necesita comprender estos mecanismos para usar TensorFlow.js

Debe sentirse cómodo (o dispuesto a intentarlo) leyendo el código fuente de TensorFlow.js para aprovechar al máximo esta guía.

Terminología

Para esta guía, es útil describir algunos términos clave por adelantado.

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 específicas de hardware/plataforma. 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.

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 degradados son código de "alto nivel" (no específicos del backend) y pueden llamar a otras operaciones o kernels.

Registro del kernel : un mapa de una tupla (nombre del kernel, nombre del backend) a una implementación del kernel.

Registro de degradado: un mapa de un nombre de kernel a una implementación de degradado .

Organización del código

Las operaciones y los gradientes se definen en tfjs-core .

Los núcleos son específicos de backend y se definen en sus respectivas carpetas de backend (p. ej ., tfjs-backend-cpu ).

No es necesario definir operaciones, kernels y gradientes personalizados dentro de estos paquetes. Pero a menudo usará 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 una salida de tensor, a menudo con tensores como entrada.

  • Algunas operaciones se pueden definir completamente en términos de operaciones existentes, y solo deben importar y llamar a estas funciones directamente. Aquí hay un ejemplo .
  • La implementación de una operación también puede enviarse a kernels específicos de backend. Esto se hace a través Engine.runKernel y se describirá con más detalle en la sección "implementación de kernels personalizados". Aquí hay un ejemplo .

Implementación de núcleos personalizados

Las implementaciones de kernel específicas de back-end permiten una implementación optimizada de la lógica para una operación determinada. Los núcleos son invocados por operaciones que llaman a tf.engine().runKernel() . Las implementaciones de un núcleo se definen por cuatro cosas

  • Un nombre de núcleo.
  • El backend en el que se implementa el kernel.
  • Entradas: Argumentos de tensor para la función kernel.
  • Atributos: Argumentos no tensoriales para la función kernel.

Aquí hay un ejemplo de una implementación del núcleo . Las convenciones que se usan para implementar son específicas del backend y se entienden mejor al observar la implementación y la documentación de cada backend en particular.

Generalmente, los núcleos operan a un nivel más bajo que los tensores y, en su lugar, leen y escriben directamente en la memoria que tfjs-core eventualmente envolverá en tensores.

Una vez que se implementa un kernel, se puede registrar con TensorFlow.js usando 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 enviar a la implementación en el backend activo actual.

Implementación de degradados personalizados

Los gradientes generalmente se definen para un kernel determinado (identificado por el mismo nombre de kernel utilizado en una llamada a tf.engine().runKernel(...) ). Esto permite que tfjs-core use un registro para buscar definiciones de gradiente para cualquier kernel en tiempo de ejecución.

La implementación de degradados personalizados es útil para:

  • Agregar una definición de gradiente que puede no estar presente en la biblioteca
  • Anular una definición de gradiente existente para personalizar el cálculo de gradiente para un kernel determinado.

Puede ver ejemplos de implementaciones de degradado aquí .

Una vez que haya implementado un gradiente para una llamada determinada, puede registrarse con TensorFlow.js usando la función registerGradient de tfjs-core.

El otro enfoque para implementar gradientes personalizados que pasa por alto el registro de gradientes (y, por lo tanto, permite calcular gradientes para funciones arbitrarias de formas arbitrarias es usar tf.customGrad .

Aquí hay un ejemplo de una operación dentro de la biblioteca de usar customGrad