Panoramica
Questa guida descrive i meccanismi per definire operazioni personalizzate (ops), kernel e gradienti in TensorFlow.js. Ha lo scopo di fornire una panoramica dei concetti principali e dei puntatori al codice che dimostrano i concetti in azione.
A chi è rivolta questa guida?
Questa è una guida abbastanza avanzata che tocca alcuni aspetti interni di TensorFlow.js, potrebbe essere particolarmente utile per i seguenti gruppi di persone:
- Utenti avanzati di TensorFlow.js interessati a personalizzare il comportamento di varie operazioni matematiche (ad esempio, ricercatori che sovrascrivono le implementazioni dei gradienti esistenti o utenti che necessitano di patch per funzionalità mancanti nella libreria)
- Utenti che creano librerie che estendono TensorFlow.js (ad esempio una libreria di algebra lineare generale costruita sulle primitive di TensorFlow.js o un nuovo backend TensorFlow.js).
- Utenti interessati a contribuire con nuove operazioni a tensorflow.js che desiderano ottenere una panoramica generale di come funzionano questi meccanismi.
Questa non è una guida all'uso generale di TensorFlow.js poiché approfondisce i meccanismi di implementazione interni. Non è necessario comprendere questi meccanismi per utilizzare TensorFlow.js
È necessario sentirsi a proprio agio (o essere disposti a provare) a leggere il codice sorgente di TensorFlow.js per sfruttare al meglio questa guida.
Terminologia
Per questa guida è utile descrivere in anticipo alcuni termini chiave.
Operazioni (Ops) — Un'operazione matematica su uno o più tensori che produce uno o più tensori come output. Le operazioni sono codice di "alto livello" e possono utilizzare altre operazioni per definire la propria logica.
Kernel : un'implementazione specifica di un'operazione legata a funzionalità hardware/piattaforma specifiche. I kernel sono di "basso livello" e specifici del backend. Alcune operazioni hanno una mappatura uno a uno dall'operazione al kernel mentre altre operazioni utilizzano più kernel.
Gradient / GradFunc — La definizione di "modalità all'indietro" di un operazione/kernel che calcola la derivata di quella funzione rispetto ad alcuni input. I gradienti sono codice di "alto livello" (non specifico del backend) e possono chiamare altre operazioni o kernel.
Registro del kernel : una mappa da una tupla (nome del kernel, nome del backend) a un'implementazione del kernel.
Registro dei gradienti : una mappa da un nome kernel a un'implementazione del gradiente .
Organizzazione del codice
Operazioni e gradienti sono definiti in tfjs-core .
I kernel sono specifici del backend e sono definiti nelle rispettive cartelle del backend (ad esempio tfjs-backend-cpu ).
Non è necessario definire operazioni personalizzate, kernel e gradienti all'interno di questi pacchetti. Ma spesso utilizzeranno simboli simili nella loro implementazione.
Implementazione di operazioni personalizzate
Un modo di pensare a un'operazione personalizzata è proprio come una funzione JavaScript che restituisce un output di tensore, spesso con tensori come input.
- Alcune operazioni possono essere completamente definite in termini di operazioni esistenti e dovrebbero semplicemente importare e chiamare direttamente queste funzioni. Ecco un esempio .
- L'implementazione di un'operazione può anche essere inviata a kernel specifici del backend. Questo viene fatto tramite
Engine.runKernele sarà descritto più avanti nella sezione "Implementazione di kernel personalizzati". Ecco un esempio .
Implementazione di kernel personalizzati
Le implementazioni del kernel specifiche del backend consentono l'implementazione ottimizzata della logica per una determinata operazione. I kernel vengono richiamati dalle operazioni che chiamano tf.engine().runKernel() . Le implementazioni del kernel sono definite da quattro cose
- Un nome del kernel.
- Il backend in cui è implementato il kernel.
- Input: argomenti tensoriali per la funzione del kernel.
- Attributi: argomenti non tensoriali della funzione kernel.
Ecco un esempio di implementazione del kernel . Le convenzioni utilizzate per l'implementazione sono specifiche del backend e si comprendono meglio osservando l'implementazione e la documentazione di ogni particolare backend.
Generalmente i kernel operano a un livello inferiore rispetto ai tensori e invece leggono e scrivono direttamente nella memoria che verrà eventualmente racchiusa nei tensori da tfjs-core.
Una volta implementato, il kernel può essere registrato con TensorFlow.js utilizzando la funzione registerKernel di tfjs-core. Puoi registrare un kernel per ogni backend in cui desideri che funzioni. Una volta registrato, il kernel può essere richiamato con tf.engine().runKernel(...) e TensorFlow.js si assicurerà di inviarlo all'implementazione nel file backend attualmente attivo.
Implementazione di gradienti personalizzati
I gradienti sono generalmente definiti per un dato kernel (identificato dallo stesso nome kernel utilizzato in una chiamata a tf.engine().runKernel(...) ). Ciò consente a tfjs-core di utilizzare un registro per cercare le definizioni di gradiente per qualsiasi kernel in fase di runtime.
L'implementazione di gradienti personalizzati è utile per:
- Aggiunta di una definizione di gradiente che potrebbe non essere presente nella libreria
- Sostituire una definizione di gradiente esistente per personalizzare il calcolo del gradiente per un determinato kernel.
Puoi vedere esempi di implementazioni del gradiente qui .
Una volta implementato un gradiente per una determinata chiamata, è possibile registrarlo con TensorFlow.js utilizzando la funzione registerGradient di tfjs-core.
L'altro approccio all'implementazione dei gradienti personalizzati che ignora il registro dei gradienti (e quindi consente di calcolare i gradienti per funzioni arbitrarie in modi arbitrari) utilizza tf.customGrad .
Ecco un esempio di un'operazione all'interno della libreria sull'utilizzo di customGrad