Visão geral
Este guia descreve os mecanismos para definir operações personalizadas (ops), kernels e gradientes no TensorFlow.js. O objetivo é fornecer uma visão geral dos principais conceitos e indicadores de código que demonstrem os conceitos em ação.
Para quem é este guia?
Este é um guia bastante avançado que aborda alguns aspectos internos do TensorFlow.js. Pode ser particularmente útil para os seguintes grupos de pessoas:
- Usuários avançados do TensorFlow.js interessados em personalizar o comportamento de diversas 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 construída sobre 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 entra em 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 uma operação vinculada a recursos específicos de hardware/plataforma. Os kernels são de 'baixo nível' e específicos de backend. Algumas operações têm um mapeamento individual de operação para kernel, enquanto outras operações usam vários kernels.
Gradient / GradFunc — A definição de 'modo retroativo' de um op/kernel que calcula a derivada dessa função em relação a alguma entrada. 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ão símbolos semelhantes em sua implementação.
Implementando operações personalizadas
Uma maneira de pensar em uma operação personalizada é 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 uma operação também pode ser enviada para kernels específicos de back-end. Isso é feito via
Engine.runKernele será descrito mais detalhadamente na seção “implementando kernels personalizados”. Aqui está um exemplo .
Implementando Kernels Personalizados
Implementações específicas de backend do kernel permitem a implementação otimizada da lógica para uma determinada operação. Kernels são invocados por operações que chamam 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 tensores para a função do kernel.
- Atributos: Argumentos não tensores para a função kernel.
Aqui está um exemplo de implementação de kernel . As convenções usadas para implementação são específicas de back-end e são melhor compreendidas observando a implementação e a documentação de cada back-end específico.
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 agrupada 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 esse kernel funcione. Uma vez registrado, o kernel pode ser invocado com tf.engine().runKernel(...) e o TensorFlow.js garantirá o envio 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 o cálculo do 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 gradiente (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 do customGrad