Escrevendo operações, kernels e gradientes personalizados no TensorFlow.js,Escrevendo operações, kernels e gradientes personalizados no TensorFlow.js

Visão geral

Este guia descreve os mecanismos para definir operações personalizadas (ops), kernels e gradientes no TensorFlow.js. Ele visa fornecer uma visão geral dos principais conceitos e indicadores de código que demonstram os conceitos em ação.

Para quem é este guia?

Este é um guia bastante avançado que aborda alguns aspectos internos do TensorFlow.js e pode ser particularmente útil para os seguintes grupos de pessoas:

  • Usuários avançados do TensorFlow.js interessados ​​em personalizar o comportamento de várias operações matemáticas (por exemplo, pesquisadores substituindo implementações de gradiente existentes ou usuários que precisam corrigir funcionalidades ausentes na biblioteca)
  • Usuários que criam bibliotecas que estendem o TensorFlow.js (por exemplo, uma biblioteca de álgebra linear geral criada com base nas primitivas do TensorFlow.js ou um novo back-end do TensorFlow.js).
  • Usuários interessados ​​em contribuir com novas operações para tensorflow.js que desejam obter uma visão geral de como esses mecanismos funcionam.

Este não é um guia para o uso geral do TensorFlow.js, pois ele aborda os mecanismos internos de implementação. Você não precisa entender esses mecanismos para usar o TensorFlow.js

Você precisa estar confortável (ou disposto a tentar) ler o código-fonte do TensorFlow.js para aproveitar ao máximo este guia.

Terminologia

Para este guia, alguns termos-chave são úteis para descrever antecipadamente.

Operações (Ops) — Uma operação matemática em um ou mais tensores que produz um ou mais tensores como saída. As operações são códigos de 'alto nível' e podem usar outras operações para definir sua lógica.

Kernel — Uma implementação específica de um op vinculado a recursos específicos de hardware/plataforma. Os kernels são de 'baixo nível' e específicos de back-end. Algumas operações têm um mapeamento um-para-um de op para kernel, enquanto outras operações usam vários kernels.

Gradient / GradFunc — A definição de 'modo inverso' de um op/kernel que calcula a derivada dessa função em relação a alguma entrada. Os gradientes são códigos de 'alto nível' (não específicos de back-end) e podem chamar outras operações ou kernels.

Registro do Kernel - Um mapa de uma tupla (nome do kernel, nome do backend) para uma implementação do kernel.

Gradient Registry — Um mapa de um nome de kernel para uma implementação de gradiente .

Organização do código

Operações e gradientes são definidos em tfjs-core .

Os kernels são específicos de backend e são definidos em suas respectivas pastas de backend (por exemplo , tfjs-backend-cpu ).

Operações personalizadas, kernels e gradientes não precisam ser definidos dentro desses pacotes. Mas muitas vezes usará símbolos semelhantes em sua implementação.

Implementação de operações personalizadas

Uma maneira de pensar em uma operação personalizada é apenas como uma função JavaScript que retorna alguma saída de tensor, geralmente com tensores como entrada.

  • Algumas operações podem ser completamente definidas em termos de operações existentes e devem apenas importar e chamar essas funções diretamente. Aqui está um exemplo .
  • A implementação de um op também pode ser despachada para kernels específicos de back-end. Isso é feito via Engine.runKernel e será descrito mais adiante na seção “implementando kernels personalizados”. Aqui está um exemplo .

Implementando kernels personalizados

As implementações de kernel específicas de back-end permitem a implementação otimizada da lógica para uma determinada operação. Kernels são invocados por ops chamando tf.engine().runKernel() . Uma implementação de kernel é definida por quatro coisas

  • Um nome de kernel.
  • O back-end em que o kernel é implementado.
  • Entradas: Argumentos do tensor para a função do kernel.
  • Atributos: argumentos não tensores para a função do kernel.

Aqui está um exemplo de uma implementação de kernel . As convenções usadas para implementar são específicas do back-end e são mais bem compreendidas observando a implementação e a documentação de cada back-end em particular.

Geralmente os kernels operam em um nível inferior aos tensores e, em vez disso, leem e gravam diretamente na memória que será eventualmente envolvida em tensores pelo tfjs-core.

Depois que um kernel é implementado, ele pode ser registrado no TensorFlow.js usando a função registerKernel do tfjs-core. Você pode registrar um kernel para cada back-end em que deseja que o kernel funcione. Uma vez registrado, o kernel pode ser invocado com tf.engine().runKernel(...) e o TensorFlow.js fará o despacho para a implementação no back-end ativo atual.

Implementando gradientes personalizados

Os gradientes são geralmente definidos para um determinado kernel (identificado pelo mesmo nome de kernel usado em uma chamada para tf.engine().runKernel(...) ). Isso permite que o tfjs-core use um registro para procurar definições de gradiente para qualquer kernel em tempo de execução.

A implementação de gradientes personalizados é útil para:

  • Adicionando uma definição de gradiente que pode não estar presente na biblioteca
  • Substituindo uma definição de gradiente existente para personalizar a computação de gradiente para um determinado kernel.

Você pode ver exemplos de implementações de gradiente aqui .

Depois de implementar um gradiente para uma determinada chamada, ele pode ser registrado no TensorFlow.js usando a função registerGradient do tfjs-core.

A outra abordagem para implementar gradientes personalizados que ignora o registro de gradientes (e, portanto, permite calcular gradientes para funções arbitrárias de maneiras arbitrárias é usar tf.customGrad .

Aqui está um exemplo de uma operação dentro da biblioteca de uso de customGrad

,

Visão geral

Este guia descreve os mecanismos para definir operações personalizadas (ops), kernels e gradientes no TensorFlow.js. Ele visa fornecer uma visão geral dos principais conceitos e indicadores de código que demonstram os conceitos em ação.

Para quem é este guia?

Este é um guia bastante avançado que aborda alguns aspectos internos do TensorFlow.js e pode ser particularmente útil para os seguintes grupos de pessoas:

  • Usuários avançados do TensorFlow.js interessados ​​em personalizar o comportamento de várias operações matemáticas (por exemplo, pesquisadores substituindo implementações de gradiente existentes ou usuários que precisam corrigir funcionalidades ausentes na biblioteca)
  • Usuários que criam bibliotecas que estendem o TensorFlow.js (por exemplo, uma biblioteca de álgebra linear geral criada com base nas primitivas do TensorFlow.js ou um novo back-end do TensorFlow.js).
  • Usuários interessados ​​em contribuir com novas operações para tensorflow.js que desejam obter uma visão geral de como esses mecanismos funcionam.

Este não é um guia para o uso geral do TensorFlow.js, pois ele aborda os mecanismos internos de implementação. Você não precisa entender esses mecanismos para usar o TensorFlow.js

Você precisa estar confortável (ou disposto a tentar) ler o código-fonte do TensorFlow.js para aproveitar ao máximo este guia.

Terminologia

Para este guia, alguns termos-chave são úteis para descrever antecipadamente.

Operações (Ops) — Uma operação matemática em um ou mais tensores que produz um ou mais tensores como saída. As operações são códigos de 'alto nível' e podem usar outras operações para definir sua lógica.

Kernel — Uma implementação específica de um op vinculado a recursos específicos de hardware/plataforma. Os kernels são de 'baixo nível' e específicos de back-end. Algumas operações têm um mapeamento um-para-um de op para kernel, enquanto outras operações usam vários kernels.

Gradient / GradFunc — A definição de 'modo inverso' de um op/kernel que calcula a derivada dessa função em relação a alguma entrada. Os gradientes são códigos de 'alto nível' (não específicos de back-end) e podem chamar outras operações ou kernels.

Registro do Kernel - Um mapa de uma tupla (nome do kernel, nome do backend) para uma implementação do kernel.

Gradient Registry — Um mapa de um nome de kernel para uma implementação de gradiente .

Organização do código

Operações e gradientes são definidos em tfjs-core .

Os kernels são específicos de backend e são definidos em suas respectivas pastas de backend (por exemplo , tfjs-backend-cpu ).

Operações personalizadas, kernels e gradientes não precisam ser definidos dentro desses pacotes. Mas muitas vezes usará símbolos semelhantes em sua implementação.

Implementação de operações personalizadas

Uma maneira de pensar em uma operação personalizada é apenas como uma função JavaScript que retorna alguma saída de tensor, geralmente com tensores como entrada.

  • Algumas operações podem ser completamente definidas em termos de operações existentes e devem apenas importar e chamar essas funções diretamente. Aqui está um exemplo .
  • A implementação de um op também pode ser despachada para kernels específicos de back-end. Isso é feito via Engine.runKernel e será descrito mais adiante na seção “implementando kernels personalizados”. Aqui está um exemplo .

Implementando kernels personalizados

As implementações de kernel específicas de back-end permitem a implementação otimizada da lógica para uma determinada operação. Kernels são invocados por ops chamando tf.engine().runKernel() . Uma implementação de kernel é definida por quatro coisas

  • Um nome de kernel.
  • O back-end em que o kernel é implementado.
  • Entradas: Argumentos do tensor para a função do kernel.
  • Atributos: argumentos não tensores para a função do kernel.

Aqui está um exemplo de uma implementação de kernel . As convenções usadas para implementar são específicas do back-end e são mais bem compreendidas observando a implementação e a documentação de cada back-end em particular.

Geralmente os kernels operam em um nível inferior aos tensores e, em vez disso, leem e gravam diretamente na memória que será eventualmente envolvida em tensores pelo tfjs-core.

Depois que um kernel é implementado, ele pode ser registrado no TensorFlow.js usando a função registerKernel do tfjs-core. Você pode registrar um kernel para cada back-end em que deseja que o kernel funcione. Uma vez registrado, o kernel pode ser invocado com tf.engine().runKernel(...) e o TensorFlow.js fará o despacho para a implementação no back-end ativo atual.

Implementando gradientes personalizados

Os gradientes são geralmente definidos para um determinado kernel (identificado pelo mesmo nome de kernel usado em uma chamada para tf.engine().runKernel(...) ). Isso permite que o tfjs-core use um registro para pesquisar definições de gradiente para qualquer kernel em tempo de execução.

A implementação de gradientes personalizados é útil para:

  • Adicionando uma definição de gradiente que pode não estar presente na biblioteca
  • Substituindo uma definição de gradiente existente para personalizar a computação de gradiente para um determinado kernel.

Você pode ver exemplos de implementações de gradiente aqui .

Depois de implementar um gradiente para uma determinada chamada, ele pode ser registrado no TensorFlow.js usando a função registerGradient do tfjs-core.

A outra abordagem para implementar gradientes personalizados que ignora o registro de gradientes (e, portanto, permite calcular gradientes para funções arbitrárias de maneiras arbitrárias é usar tf.customGrad .

Aqui está um exemplo de uma operação dentro da biblioteca de uso de customGrad