Crea un'operazione

Se desideri creare un'operazione che non sia coperta dalla libreria TensorFlow esistente, ti consigliamo di provare prima a scrivere l'operazione in Python come una composizione di operazioni o funzioni Python esistenti. Se ciò non è possibile, puoi creare un'operazione C++ personalizzata. Esistono diversi motivi per cui potresti voler creare un'operazione C++ personalizzata:

  • Non è facile né possibile esprimere la propria operazione come una composizione di operazioni esistenti.
  • Non è efficiente esprimere la tua operazione come una composizione di primitive esistenti.
  • Vuoi fondere manualmente una composizione di primitive che un futuro compilatore troverebbe difficile fondere.

Ad esempio, immagina di voler implementare qualcosa come il "pooling mediano", simile all'operatore "MaxPool", ma calcolando le mediane su finestre scorrevoli invece che su valori massimi. Potrebbe essere possibile farlo utilizzando una composizione di operazioni (ad esempio, utilizzando ExtractImagePatches e TopK), ma potrebbe non essere efficiente in termini di prestazioni o memoria come un'operazione nativa in cui è possibile fare qualcosa di più intelligente in un'unica operazione fusa. Come sempre, in genere vale prima la pena provare a esprimere ciò che si desidera utilizzando la composizione dell'operatore, scegliendo di aggiungere una nuova operazione solo se ciò si rivela difficile o inefficiente.

Per incorporare la tua operazione personalizzata dovrai:

  1. Registra la nuova operazione in un file C++. La registrazione dell'operazione definisce un'interfaccia (specifica) per la funzionalità dell'operazione, che è indipendente dall'implementazione dell'operazione. Ad esempio, la registrazione dell'operazione definisce il nome dell'operazione e gli input e gli output dell'operazione. Definisce inoltre la funzione di forma utilizzata per l'inferenza della forma del tensore.
  2. Implementare l'operazione in C++. L'implementazione di un'operazione è nota come kernel ed è l'implementazione concreta della specifica registrata nel passaggio 1. Possono esserci più kernel per diversi tipi o architetture di input/output (ad esempio CPU, GPU).
  3. Crea un wrapper Python (facoltativo). Questo wrapper è l'API pubblica utilizzata per creare l'operazione in Python. Dalla registrazione dell'operazione viene generato un wrapper predefinito che può essere utilizzato direttamente o aggiunto.
  4. Scrivi una funzione per calcolare i gradienti per l'operazione (opzionale).
  5. Prova l'operazione. Di solito lo facciamo in Python per comodità, ma puoi anche testare l'operazione in C++. Se definisci i gradienti, puoi verificarli con Python tf.test.compute_gradient_error . Vedi relu_op_test.py come esempio che testa le funzioni forward degli operatori simili a Relu e i loro gradienti.

Prerequisiti

Definire l'interfaccia operativa

L'interfaccia di un'operazione viene definita dall'utente registrandola nel sistema TensorFlow. Nella registrazione, specifichi il nome della tua operazione, i suoi input (tipi e nomi) e output (tipi e nomi), nonché le docstring e qualsiasi attributo che l'operazione potrebbe richiedere.

Per vedere come funziona, supponiamo che desideri creare un op che accetta un tensore di int32 s e restituisce una copia del tensore, con tutti gli elementi tranne il primo impostati su zero. Per fare ciò, crea un file denominato zero_out.cc . Quindi aggiungi una chiamata alla macro REGISTER_OP che definisce l'interfaccia per la tua operazione:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

Questa operazione ZeroOut accetta un tensore to_zero di interi a 32 bit come input e restituisce un tensore zeroed di interi a 32 bit. L'operazione utilizza anche una funzione di forma per garantire che il tensore di output abbia la stessa forma del tensore di input. Ad esempio, se l'input è un tensore di forma [10, 20], allora questa funzione di forma specifica che anche la forma di output è [10, 20].

Implementare il kernel per l'op

Dopo aver definito l'interfaccia, fornire una o più implementazioni dell'op. Per creare uno di questi kernel, crea una classe che estenda OpKernel e sovrascriva il metodo Compute . Il metodo Compute fornisce un argomento context di tipo OpKernelContext* , da cui è possibile accedere a elementi utili come i tensori di input e output.

Aggiungi il tuo kernel al file che hai creato sopra. Il kernel potrebbe assomigliare a questo:

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

Dopo aver implementato il tuo kernel, lo registri con il sistema TensorFlow. Nella registrazione si specificano diversi vincoli in base ai quali verrà eseguito questo kernel. Ad esempio, potresti avere un kernel realizzato per le CPU e uno separato per le GPU.

Per fare ciò per l'operazione ZeroOut , aggiungi quanto segue a zero_out.cc :

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

Kernel CPU multi-thread

Per scrivere un kernel CPU multi-thread, è possibile utilizzare la funzione Shard in work_sharder.h . Questa funzione suddivide una funzione di calcolo tra i thread configurati per essere utilizzata per il threading intra-op (vedere intra_op_parallelism_threads in config.proto ).

Kernel GPU

Un kernel GPU è implementato in due parti: OpKernel e il kernel CUDA e il suo codice di lancio.

A volte l'implementazione OpKernel è comune tra un kernel CPU e GPU, ad esempio per quanto riguarda l'ispezione degli input e l'allocazione degli output. In tal caso, un'implementazione suggerita è quella di:

  1. Definire il modello OpKernel sul dispositivo e il tipo primitivo del tensore.
  2. Per eseguire il calcolo effettivo dell'output, la funzione Compute chiama una struttura funtore basata su modello.
  3. La specializzazione di quel funtore per CPUDevice è definita nello stesso file, ma la specializzazione per GPUDevice è definita in un file .cu.cc, poiché verrà compilato con il compilatore CUDA.

Ecco un esempio di implementazione. 

// kernel_example.h
#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

#include <unsupported/Eigen/CXX11/Tensor>

template <typename Device, typename T>
struct ExampleFunctor {
  void operator()(const Device& d, int size, const T* in, T* out);
};

#if GOOGLE_CUDA
// Partially specialize functor for GpuDevice.
template <typename T>
struct ExampleFunctor<Eigen::GpuDevice, T> {
  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);
};
#endif

#endif KERNEL_EXAMPLE_H_

// kernel_example.cc
#include "kernel_example.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

REGISTER_OP("Example")
    .Attr("T: numbertype")
    .Input("input: T")
    .Output("input_times_two: T")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

// CPU specialization of actual computation.
template <typename T>
struct ExampleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, int size, const T* in, T* out) {
    for (int i = 0; i < size; ++i) {
      out[i] = 2 * in[i];
    }
  }
};

// OpKernel definition.
// template parameter <T> is the datatype of the tensors.
template <typename Device, typename T>
class ExampleOp : public OpKernel {
 public:
  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));

    // Do the computation.
    OP_REQUIRES(context, input_tensor.NumElements() <= tensorflow::kint32max,
                errors::InvalidArgument("Too many elements in tensor"));
    ExampleFunctor<Device, T>()(
        context->eigen_device<Device>(),
        static_cast<int>(input_tensor.NumElements()),
        input_tensor.flat<T>().data(),
        output_tensor->flat<T>().data());
  }
};

// Register the CPU kernels.
#define REGISTER_CPU(T)                                          \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
      ExampleOp<CPUDevice, T>);
REGISTER_CPU(float);
REGISTER_CPU(int32);

// Register the GPU kernels.
#ifdef GOOGLE_CUDA
#define REGISTER_GPU(T)                                          \
  /* Declare explicit instantiations in kernel_example.cu.cc. */ \
  extern template class ExampleFunctor<GPUDevice, T>;            \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
      ExampleOp<GPUDevice, T>);
REGISTER_GPU(float);
REGISTER_GPU(int32);
#endif  // GOOGLE_CUDA

// kernel_example.cu.cc
#ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "kernel_example.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

using namespace tensorflow;

using GPUDevice = Eigen::GpuDevice;

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    out[i] = 2 * __ldg(in + i);
  }
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const GPUDevice& d, int size, const T* in, T* out) {
  // Launch the cuda kernel.
  //
  // See core/util/gpu_kernel_helper.h for example of computing
  // block count and thread_per_block count.
  int block_count = 1024;
  int thread_per_block = 20;
  ExampleCudaKernel<T>
      <<<block_count, thread_per_block, 0, d.stream()>>>(size, in, out);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, int32>;

#endif  // GOOGLE_CUDA

Costruisci la libreria operativa

Compila l'operazione utilizzando il compilatore di sistema (installazione binaria TensorFlow)

Dovresti essere in grado di compilare zero_out.cc con un compilatore C++ come g++ o clang disponibile sul tuo sistema. Il pacchetto PIP binario installa i file di intestazione e la libreria necessari per compilare l'operazione in posizioni specifiche del sistema. Tuttavia, la libreria Python TensorFlow fornisce la funzione get_include per ottenere la directory dell'intestazione e la directory get_lib ha un oggetto condiviso a cui collegarsi. Ecco gli output di queste funzioni su una macchina Ubuntu.

$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'/usr/local/lib/python3.6/site-packages/tensorflow/include'
>>> tf.sysconfig.get_lib()
'/usr/local/lib/python3.6/site-packages/tensorflow'

Supponendo che tu abbia installato g++ , ecco la sequenza di comandi che puoi utilizzare per compilare la tua operazione in una libreria dinamica.

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++14 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

Su macOS, è richiesto il flag aggiuntivo "-unfine Dynamic_lookup" durante la creazione del file .so .

Nota sulla versione gcc >=5 : gcc utilizza la nuova ABI C++ a partire dalla versione 5 . TensorFlow 2.8 e versioni precedenti sono state realizzate con gcc4 che utilizza la vecchia ABI. Se stai utilizzando queste versioni di TensorFlow e stai tentando di compilare la tua libreria op con gcc>=5 , aggiungi -D_GLIBCXX_USE_CXX11_ABI=0 alla riga di comando per rendere la libreria compatibile con la versione ABI precedente. I pacchetti TensorFlow 2.9+ sono compatibili con la nuova ABI per impostazione predefinita.

Compila l'operazione utilizzando bazel (installazione del sorgente TensorFlow)

Se hai installato i sorgenti di TensorFlow, puoi utilizzare il sistema di compilazione di TensorFlow per compilare il tuo op. Inserisci un file BUILD con la seguente regola di compilazione Bazel nella directory tensorflow/core/user_ops .

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

Esegui il comando seguente per creare zero_out.so .

$ bazel build --config opt //tensorflow/core/user_ops:zero_out.so

Per compilare l'operazione Example , con il kernel CUDA, è necessario utilizzare il parametro gpu_srcs di tf_custom_op_library . Inserisci un file BUILD con la seguente regola di compilazione Bazel in una nuova cartella all'interno della directory tensorflow/core/user_ops (ad esempio "example_gpu").

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    # kernel_example.cc  kernel_example.cu.cc  kernel_example.h
    name = "kernel_example.so",
    srcs = ["kernel_example.h", "kernel_example.cc"],
    gpu_srcs = ["kernel_example.cu.cc", "kernel_example.h"],
)

Esegui il comando seguente per compilare kernel_example.so .

$ bazel build --config opt //tensorflow/core/user_ops/example_gpu:kernel_example.so

Usa l'operazione in Python

L'API TensorFlow Python fornisce la funzione tf.load_op_library per caricare la libreria dinamica e registrare l'operazione con il framework TensorFlow. load_op_library restituisce un modulo Python che contiene i wrapper Python per l'operazione e il kernel. Pertanto, una volta creato l'operazione, puoi eseguire le seguenti operazioni da Python:

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
print(zero_out_module.zero_out([[1, 2], [3, 4]]).numpy())

# Prints
array([[1, 0], [0, 0]], dtype=int32)

Tieni presente che alla funzione generata verrà assegnato un nome snake_case (per conformarsi a PEP8 ). Pertanto, se la tua operazione è denominata ZeroOut nei file C++, la funzione Python verrà chiamata zero_out .

Per rendere l'operazione disponibile come una normale funzione import -able da un modulo Python, potrebbe essere utile avere la chiamata load_op_library in un file sorgente Python come segue: 

import tensorflow as tf

zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out

Verificare che l'operazione funzioni

Un buon modo per verificare di aver implementato con successo la tua operazione è scrivere un test per essa. Crea il file zero_out_op_test.py con il contenuto:

import tensorflow as tf

class ZeroOutTest(tf.test.TestCase):
  def testZeroOut(self):
    zero_out_module = tf.load_op_library('./zero_out.so')
    with self.test_session():
      result = zero_out_module.zero_out([5, 4, 3, 2, 1])
      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

if __name__ == "__main__":
  tf.test.main()

Quindi esegui il test (supponendo che tu abbia installato tensorflow): 

$ python zero_out_op_test.py

Integra funzionalità avanzate nella tua operazione

Ora che sai come creare un'operazione e un'implementazione di base (e in qualche modo limitata), esamineremo alcune delle cose più complicate che in genere dovrai integrare nella tua operazione. Ciò include:

Verifiche condizionali e validazione

L'esempio sopra presuppone che l'operazione si applichi a un tensore di qualsiasi forma. E se si applicasse solo ai vettori? Ciò significa aggiungere un controllo all'implementazione OpKernel di cui sopra.

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
                errors::InvalidArgument("ZeroOut expects a 1-D vector."));
    // ...
  }

Ciò asserisce che l'input è un vettore e restituisce di aver impostato lo stato InvalidArgument se non lo è. La macro OP_REQUIRES accetta tre argomenti:

In alternativa, se vuoi verificare se un oggetto Status restituito da qualche funzione è un errore e, in tal caso, restituirlo, usa OP_REQUIRES_OK . Entrambe queste macro ritornano dalla funzione in caso di errore.

Registrazione op

Attr

Le operazioni possono avere attributi, i cui valori vengono impostati quando l'operazione viene aggiunta a un grafico. Questi vengono utilizzati per configurare l'operazione e ai loro valori è possibile accedere sia all'interno dell'implementazione del kernel che nei tipi di input e output nella registrazione dell'operazione. Preferisci utilizzare un input anziché un attr quando possibile, poiché gli input sono più flessibili. Questo perché gli attr sono costanti e devono essere definiti al momento della costruzione del grafico. Al contrario, gli input sono tensori i cui valori possono essere dinamici; cioè, gli input possono cambiare ogni passaggio, essere impostati utilizzando un feed, ecc. Gli attributi vengono utilizzati per cose che non possono essere fatte con gli input: qualsiasi configurazione che influenzi la firma (numero o tipo di input o output) o che può' t cambiare passo dopo passo.

Definisci un attr quando registri l'operazione, specificandone il nome e il tipo utilizzando il metodo Attr , che prevede una specifica del modulo:

<name>: <attr-type-expr>

dove <name> inizia con una lettera e può essere composto da caratteri alfanumerici e trattini bassi e <attr-type-expr> è un'espressione di tipo nella forma descritta di seguito .

Ad esempio, se desideri che l'operazione ZeroOut conservi un indice specificato dall'utente, anziché solo l'elemento 0, puoi registrare l'operazione in questo modo:

REGISTER_OP("ZeroOut")
    .Attr("preserve_index: int")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

(Si noti che l'insieme dei tipi di attributi è diverso dal tf.DType utilizzato per input e output.)

Il tuo kernel può quindi accedere a questo attr nel suo costruttore tramite il parametro context :

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
    // Get the index of the value to preserve
    OP_REQUIRES_OK(context,
                   context->GetAttr("preserve_index", &preserve_index_));
    // Check that preserve_index is positive
    OP_REQUIRES(context, preserve_index_ >= 0,
                errors::InvalidArgument("Need preserve_index >= 0, got ",
                                        preserve_index_));
  }
  void Compute(OpKernelContext* context) override {
    // ...
  }
 private:
  int preserve_index_;
};

che può quindi essere utilizzato nel metodo Compute : 

  void Compute(OpKernelContext* context) override {
    // ...

    // We're using saved attr to validate potentially dynamic input
    // So we check that preserve_index is in range
    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
                errors::InvalidArgument("preserve_index out of range"));

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the requested input value
    output_flat(preserve_index_) = input(preserve_index_);
  }

Tipi di attr

I seguenti tipi sono supportati in un attr:

  • string : qualsiasi sequenza di byte (non è necessario che sia UTF8).
  • int : un numero intero con segno.
  • float : un numero in virgola mobile.
  • bool : vero o falso.
  • type : uno dei valori (non di riferimento) di DataType .
  • shape : A TensorShapeProto .
  • list(<type>) : un elenco di <type> , dove <type> è uno dei tipi precedenti. Tieni presente che list(list(<type>)) non è valido.

Vedi anche: op_def_builder.cc:FinalizeAttr per un elenco definitivo.

Valori predefiniti e vincoli

Gli attributi possono avere valori predefiniti e alcuni tipi di attributi possono avere vincoli. Per definire un attributo con vincoli, puoi utilizzare i seguenti <attr-type-expr> s:

{'<string1>', '<string2>'} : il valore deve essere una stringa che ha il valore <string1> o <string2> . Il nome del tipo, string , è implicito quando si utilizza questa sintassi. Questo emula un'enumerazione:

REGISTER_OP("EnumExample")
    .Attr("e: {'apple', 'orange'}");

{<type1>, <type2>} : il valore è di tipo type e deve essere uno tra <type1> o <type2> , dove <type1> e <type2> sono supportati tf.DType . Non specifichi che il tipo di attr è type . Ciò è implicito quando si dispone di un elenco di tipi in {...} . Ad esempio, in questo caso attr t è un tipo che deve essere int32 , float o bool :

REGISTER_OP("RestrictedTypeExample")
    .Attr("t: {int32, float, bool}");

Esistono scorciatoie per i vincoli di tipo comuni:

  • numbertype : type di tipo limitato ai tipi numerici (non stringa e non bool).
  • realnumbertype : Come numbertype senza tipi complessi.
  • quantizedtype : come numbertype ma solo i tipi di numeri quantizzati.

Gli elenchi specifici di tipi consentiti da questi sono definiti dalle funzioni (come NumberTypes() ) in tensorflow/core/framework/types.h . In questo esempio l'atttr t deve essere uno dei tipi numerici:

REGISTER_OP("NumberType")
    .Attr("t: numbertype");

Per questa operazione:

tf.number_type(t=tf.int32)  # Valid
tf.number_type(t=tf.bool)   # Invalid

Gli elenchi possono essere combinati con altri elenchi e tipologie singole. L'operazione seguente consente ad attr t di essere uno qualsiasi dei tipi numerici o il tipo bool:

REGISTER_OP("NumberOrBooleanType")
    .Attr("t: {numbertype, bool}");

Per questa operazione:

tf.number_or_boolean_type(t=tf.int32)  # Valid
tf.number_or_boolean_type(t=tf.bool)   # Valid
tf.number_or_boolean_type(t=tf.string) # Invalid

int >= <n> : il valore deve essere un int il cui valore è maggiore o uguale a <n> , dove <n> è un numero naturale. Ad esempio, la seguente registrazione op specifica che attr a deve avere un valore almeno 2 :

REGISTER_OP("MinIntExample")
    .Attr("a: int >= 2");

list(<type>) >= <n> : un elenco di tipo <type> la cui lunghezza è maggiore o uguale a <n> . Ad esempio, la seguente registrazione op specifica che attr a è un elenco di tipi ( int32 o float ) e che devono essercene almeno 3:

REGISTER_OP("TypeListExample")
    .Attr("a: list({int32, float}) >= 3");

Per impostare un valore predefinito per un attr (rendendolo facoltativo nel codice generato), aggiungi = <default> alla fine, come in:

REGISTER_OP("AttrDefaultExample")
    .Attr("i: int = 0");

Inoltre, è possibile specificare sia un vincolo che un valore predefinito:

REGISTER_OP("AttrConstraintAndDefaultExample")
    .Attr("i: int >= 1 = 1");

La sintassi supportata del valore predefinito è quella che verrebbe utilizzata nella rappresentazione prototipale della definizione GraphDef risultante.

Ecco alcuni esempi su come specificare un valore predefinito per tutti i tipi:

REGISTER_OP("AttrDefaultExampleForAllTypes")
   .Attr("s: string = 'foo'")
   .Attr("i: int = 0")
   .Attr("f: float = 1.0")
   .Attr("b: bool = true")
   .Attr("ty: type = DT_INT32")
   .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
   .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
   .Attr("l_empty: list(int) = []")
   .Attr("l_int: list(int) = [2, 3, 5, 7]");

Si noti in particolare che i valori di tipo type utilizzano tf.DType .

Polimorfismo

Polimorfismo del tipo

Per le operazioni che possono accettare tipi diversi come input o produrre diversi tipi di output, puoi specificare un attributo in un tipo di input o output nella registrazione dell'operazione. In genere dovresti quindi registrare un OpKernel per ciascun tipo supportato.

Ad esempio, se desideri che l'operazione ZeroOut funzioni su float oltre che su int32 , la registrazione dell'operazione potrebbe essere simile a:

REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

La registrazione dell'operazione ora specifica che il tipo di input deve essere float , o int32 , e che il suo output sarà dello stesso tipo, poiché entrambi hanno il tipo T .

Denominazione

A input, output e attributi generalmente dovrebbero essere assegnati nomi Snake_case. L'unica eccezione sono gli attributi utilizzati come tipo di input o nel tipo di output. Tali attributi possono essere dedotti quando l'operazione viene aggiunta al grafico e quindi non compaiono nella funzione dell'operazione. Ad esempio, quest'ultima definizione di ZeroOut genererà una funzione Python simile a:

def zero_out(to_zero, name=None):
  """...
  Args:
    to_zero: A `Tensor`. Must be one of the following types:
        `float32`, `int32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `to_zero`.
  """

Se to_zero viene passato un tensore int32 , allora T viene automaticamente impostato su int32 (beh, in realtà DT_INT32 ). A queste attribuzioni dedotte vengono assegnati nomi in maiuscolo o CamelCase.

Confronta questo con un op che ha un tipo attr che determina il tipo di output:

REGISTER_OP("StringToNumber")
    .Input("string_tensor: string")
    .Output("output: out_type")
    .Attr("out_type: {float, int32} = DT_FLOAT");
    .Doc(R"doc(
Converts each string in the input Tensor to the specified numeric type.
)doc");

In questo caso, l'utente deve specificare il tipo di output, come nel Python generato: 

def string_to_number(string_tensor, out_type=None, name=None):
  """Converts each string in the input Tensor to the specified numeric type.

  Args:
    string_tensor: A `Tensor` of type `string`.
    out_type: An optional `tf.DType` from: `tf.float32, tf.int32`.
      Defaults to `tf.float32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor` of type `out_type`.
  """
Esempio di polimorfismo del tipo
#include "tensorflow/core/framework/op_kernel.h"

class ZeroOutInt32Op : public OpKernel {
  // as before
};

class ZeroOutFloatOp : public OpKernel {
 public:
  explicit ZeroOutFloatOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<float>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<float>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutInt32Op);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutFloatOp);

Per preservare la compatibilità con le versioni precedenti , dovresti specificare un valore predefinito quando aggiungi un attr a un'operazione esistente:

REGISTER_OP("ZeroOut")
  .Attr("T: {float, int32} = DT_INT32")
  .Input("to_zero: T")
  .Output("zeroed: T")

Diciamo che vuoi aggiungere più tipi, diciamo double :

REGISTER_OP("ZeroOut")
    .Attr("T: {float, double, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

Invece di scrivere un altro OpKernel con codice ridondante come sopra, spesso sarai in grado di utilizzare invece un modello C++. Avrai comunque una registrazione del kernel (chiamata REGISTER_KERNEL_BUILDER ) per sovraccarico.

template <typename T>
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<T>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<T>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutOp<int32>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    ZeroOutOp<double>);

Se hai più di un paio di sovraccarichi, puoi inserire la registrazione in una macro.

#include "tensorflow/core/framework/op_kernel.h"

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);

#undef REGISTER_KERNEL

A seconda dell'elenco dei tipi per cui stai registrando il kernel, potresti essere in grado di utilizzare una macro fornita da tensorflow/core/framework/register_types.h : 

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

REGISTER_OP("ZeroOut")
    .Attr("T: realnumbertype")
    .Input("to_zero: T")
    .Output("zeroed: T");

template <typename T>
class ZeroOutOp : public OpKernel { ... };

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

#undef REGISTER_KERNEL
Elenca input e output

Oltre a poter accettare o produrre tipi diversi, gli op possono consumare o produrre un numero variabile di tensori.

Nell'esempio successivo, attr T contiene un elenco di tipi e viene utilizzato come tipo sia in l'input che per l'output out . L'input e l'output sono elenchi di tensori di quel tipo (e il numero e i tipi di tensori nell'output sono gli stessi dell'input, poiché entrambi hanno il tipo T ).

REGISTER_OP("PolymorphicListExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

È inoltre possibile imporre restrizioni sui tipi che possono essere specificati nell'elenco. Nel prossimo caso, l'input è un elenco di tensori float e double . L'operazione accetta, ad esempio, tipi di input (float, double, float) e in tal caso anche il tipo di output sarebbe (float, double, float) .

REGISTER_OP("ListTypeRestrictionExample")
    .Attr("T: list({float, double})")
    .Input("in: T")
    .Output("out: T");

Se vuoi che tutti i tensori in un elenco siano dello stesso tipo, potresti fare qualcosa del tipo:

REGISTER_OP("IntListInputExample")
    .Attr("N: int")
    .Input("in: N * int32")
    .Output("out: int32");

Questo accetta un elenco di tensori int32 e utilizza un int attr N per specificare la lunghezza dell'elenco.

Anche questo può essere reso polimorfico . Nell'esempio successivo, l'input è un elenco di tensori (con lunghezza "N" ) dello stesso tipo (ma non specificato) ( "T" ) e l'output è un singolo tensore di tipo corrispondente:

REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

Per impostazione predefinita, gli elenchi tensoriali hanno una lunghezza minima pari a 1. È possibile modificare tale impostazione predefinita utilizzando un vincolo ">=" sul corrispondente attr . Nel prossimo esempio, l'input è un elenco di almeno 2 tensori int32 :

REGISTER_OP("MinLengthIntListExample")
    .Attr("N: int >= 2")
    .Input("in: N * int32")
    .Output("out: int32");

La stessa sintassi funziona con gli attributi "list(type)" : 

REGISTER_OP("MinimumLengthPolymorphicListExample")
    .Attr("T: list(type) >= 3")
    .Input("in: T")
    .Output("out: T");

Ingressi e uscite

Per riassumere quanto sopra, una registrazione operativa può avere più input e output:

REGISTER_OP("MultipleInsAndOuts")
    .Input("y: int32")
    .Input("z: float")
    .Output("a: string")
    .Output("b: int32");

Ciascuna specifica di input o output ha la forma:

<name>: <io-type-expr>

dove <name> inizia con una lettera e può essere composto da caratteri alfanumerici e trattini bassi. <io-type-expr> è una delle seguenti espressioni di tipo:

  • <type> , dove <type> è un tipo di input supportato (ad esempio float , int32 , string ). Questo specifica un singolo tensore del tipo specificato.

    Vedi tf.DType .

    REGISTER_OP("BuiltInTypesExample")
    .Input("integers: int32")
    .Input("complex_numbers: complex64");

  • <attr-type> , dove <attr-type> è il nome di un Attr con type o list(type) (con una possibile restrizione sul tipo). Questa sintassi consente operazioni polimorfiche .

    REGISTER_OP("PolymorphicSingleInput")
    .Attr("T: type")
    .Input("in: T");

REGISTER_OP("RestrictedPolymorphicSingleInput")
    .Attr("T: {int32, int64}")
    .Input("in: T");

    Fare riferimento a un attributo di tipo list(type) consente di accettare una sequenza di tensori.

    REGISTER_OP("ArbitraryTensorSequenceExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

REGISTER_OP("RestrictedTensorSequenceExample")
    .Attr("T: list({int32, int64})")
    .Input("in: T")
    .Output("out: T");

    Si noti che il numero e il tipo di tensori nell'output out è lo stesso dell'input in , poiché entrambi sono di tipo T .

  • Per una sequenza di tensori con lo stesso tipo: <number> * <type> , dove <number> è il nome di un Attr con tipo int . Il <type> può essere tf.DType o il nome di un attr con tipo type . Come esempio del primo, questa operazione accetta un elenco di tensori int32 :

    REGISTER_OP("Int32SequenceExample")
    .Attr("NumTensors: int")
    .Input("in: NumTensors * int32")

    Mentre questa operazione accetta un elenco di tensori di qualsiasi tipo, purché siano tutti uguali:

    REGISTER_OP("SameTypeSequenceExample")
    .Attr("NumTensors: int")
    .Attr("T: type")
    .Input("in: NumTensors * T")

  • Per un riferimento a un tensore: Ref(<type>) , dove <type> è uno dei tipi precedenti.

Verrà dedotto qualsiasi attributo utilizzato nel tipo di input. Per convenzione tali attribuzioni dedotte utilizzano nomi maiuscoli (come T o N ). Altrimenti input, output e attributi hanno nomi come parametri di funzione (ad esempio num_outputs ). Per ulteriori dettagli, vedere la sezione precedente sulla denominazione dei file .

Per ulteriori dettagli, vedere tensorflow/core/framework/op_def_builder.h .

Compatibilità con le versioni precedenti

Supponiamo che tu abbia scritto un'operazione piacevole e personalizzata e l'abbia condivisa con altri, in modo da avere clienti soddisfatti che utilizzano la tua operazione. Tuttavia, vorresti apportare modifiche all'operazione in qualche modo.

In generale, le modifiche alle specifiche esistenti e archiviate devono essere compatibili con le versioni precedenti: la modifica delle specifiche di un'operazione non deve interrompere i buffer del protocollo GraphDef serializzati precedenti costruiti da specifiche precedenti. I dettagli della compatibilità GraphDef sono descritti qui .

Esistono diversi modi per preservare la compatibilità con le versioni precedenti.

  1. Qualsiasi nuovo attributo aggiunto a un'operazione deve avere valori predefiniti definiti e con tale valore predefinito l'operazione deve avere il comportamento originale. Per modificare un'operazione da non polimorfica a polimorfica, è necessario assegnare un valore predefinito al nuovo tipo attr per preservare la firma originale per impostazione predefinita. Ad esempio, se la tua operazione fosse:

    REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: float")
    .Output("out: float");

    puoi renderlo polimorfico in modo retrocompatibile usando:

    REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: T")
    .Output("out: T")
    .Attr("T: numerictype = DT_FLOAT");

  2. Puoi tranquillamente rendere un vincolo su un attr meno restrittivo. Ad esempio, puoi cambiare da {int32, int64} a {int32, int64, float} o type . Oppure puoi cambiare da {"apple", "orange"} a {"apple", "banana", "orange"} o string .

  3. È possibile modificare singoli input/output in input/output di elenco, purché l'impostazione predefinita per il tipo di elenco corrisponda alla vecchia firma.

  4. È possibile aggiungere un nuovo elenco input/output, se per impostazione predefinita è vuoto.

  5. Dai uno spazio ai nomi a tutte le nuove operazioni che crei, anteponendo ai nomi delle operazioni qualcosa di unico per il tuo progetto. Ciò evita che la tua operazione entri in collisione con eventuali operazioni che potrebbero essere incluse nelle versioni future di TensorFlow.

  6. Pianifica in anticipo! Prova ad anticipare gli usi futuri dell'op. Alcune modifiche alla firma non possono essere eseguite in modo compatibile (ad esempio, creando un elenco dello stesso tipo in un elenco di tipi diversi).

L'elenco completo delle modifiche sicure e non sicure è disponibile in tensorflow/core/framework/op_compatibility_test.cc . Se non è possibile rendere compatibile con le versioni precedenti la modifica a un'operazione, creare una nuova operazione con un nuovo nome e con la nuova semantica.

Tieni inoltre presente che, sebbene queste modifiche possano mantenere la compatibilità GraphDef , il codice Python generato potrebbe cambiare in un modo non compatibile con i vecchi chiamanti. L'API Python può essere mantenuta compatibile mediante attente modifiche in un wrapper Python scritto a mano, mantenendo la vecchia firma tranne eventualmente aggiungendo nuovi argomenti opzionali alla fine. Generalmente le modifiche incompatibili possono essere apportate solo quando TensorFlow modifica le versioni principali e devono essere conformi alla semantica della versione GraphDef .

Supporto GPU

Puoi implementare diversi OpKernel e registrarne uno per la CPU e un altro per la GPU, proprio come puoi registrare kernel per diversi tipi . Esistono diversi esempi di kernel con supporto GPU in tensorflow/core/kernels/ . Notare che alcuni kernel hanno una versione della CPU in un file .cc , una versione della GPU in un file che termina con _gpu.cu.cc e parte del codice condiviso in comune in un file .h .

Ad esempio, tf.pad ha tutto tranne il kernel GPU in tensorflow/core/kernels/pad_op.cc . Il kernel GPU è in tensorflow/core/kernels/pad_op_gpu.cu.cc e il codice condiviso è una classe basata su modelli definita in tensorflow/core/kernels/pad_op.h . Organizziamo il codice in questo modo per due motivi: consente di condividere codice comune tra le implementazioni CPU e GPU e inserisce l'implementazione GPU in un file separato in modo che possa essere compilato solo dal compilatore GPU.

Una cosa da notare, anche quando viene utilizzata la versione del kernel GPU del pad , ha comunque bisogno del suo input "paddings" nella memoria della CPU. Per contrassegnare che gli input o gli output vengono mantenuti sulla CPU, aggiungi una chiamata HostMemory() alla registrazione del kernel, ad esempio: 

#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint<T>("T")  \
                              .HostMemory("paddings"), \
                          PadOp<GPUDevice, T>)

Compilazione del kernel per il dispositivo GPU

Guarda cuda_op_kernel.cu.cc per un esempio che utilizza un kernel CUDA per implementare un'operazione. tf_custom_op_library accetta un argomento gpu_srcs in cui è possibile specificare l'elenco dei file di origine contenenti i kernel CUDA (file *.cu.cc ). Per l'utilizzo con un'installazione binaria di TensorFlow, i kernel CUDA devono essere compilati con il compilatore nvcc di NVIDIA. Ecco la sequenza di comandi che puoi utilizzare per compilare cuda_op_kernel.cu.cc e cuda_op_kernel.cc in un'unica libreria caricabile dinamicamente:

nvcc -std=c++14 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++14 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

cuda_op_kernel.so prodotto sopra può essere caricato come al solito in Python, utilizzando la funzione tf.load_op_library .

Tieni presente che se le tue librerie CUDA non sono installate in /usr/local/lib64 , dovrai specificare il percorso esplicitamente nel secondo comando (g++) sopra. Ad esempio, aggiungi -L /usr/local/cuda-8.0/lib64/ se CUDA è installato in /usr/local/cuda-8.0 .

Implementa il gradiente in Python

Dato un grafico di operazioni, TensorFlow utilizza la differenziazione automatica (backpropagation) per aggiungere nuove operazioni che rappresentano gradienti rispetto alle operazioni esistenti. Per far funzionare la differenziazione automatica per le nuove operazioni, è necessario registrare una funzione gradiente che calcola i gradienti rispetto agli input delle operazioni dati i gradienti rispetto agli output delle operazioni.

Matematicamente, se un'operazione calcola \(y = f(x)\) l'operazione gradiente registrata converte i gradienti \(\partial L/ \partial y\) di perdita \(L\) riguardo a\(y\) in gradienti \(\partial L/ \partial x\) riguardo a \(x\) tramite la regola della catena:

\[\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial f}{\partial x}.\]

Nel caso di ZeroOut , solo una voce nell'input influisce sull'output, quindi il gradiente rispetto all'input è un tensore sparso "one hot". Ciò è espresso come segue:

from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
  """The gradients for `zero_out`.

  Args:
    op: The `zero_out` `Operation` that we are differentiating, which we can use
      to find the inputs and outputs of the original op.
    grad: Gradient with respect to the output of the `zero_out` op.

  Returns:
    Gradients with respect to the input of `zero_out`.
  """
  to_zero = op.inputs[0]
  shape = array_ops.shape(to_zero)
  index = array_ops.zeros_like(shape)
  first_grad = array_ops.reshape(grad, [-1])[0]
  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
  return [to_zero_grad]  # List of one Tensor, since we have one input

Dettagli sulla registrazione delle funzioni gradiente con tf.RegisterGradient :

  • Per un'operazione con un output, la funzione gradiente prenderà un tf.Operation , op , e un tf.Tensor grad e creerà nuove operazioni dai tensori op.inputs[i] , op.outputs[i] e grad . Le informazioni su eventuali attributi possono essere trovate tramite tf.Operation.get_attr .

  • Se l'operazione ha più output, la funzione gradiente prenderà op e grads , dove grads è un elenco di gradienti rispetto a ciascun output. Il risultato della funzione gradiente deve essere un elenco di oggetti Tensor che rappresentano i gradienti rispetto a ciascun input.

  • Se non esiste un gradiente ben definito per alcuni input, come per gli input interi utilizzati come indici, il gradiente restituito corrispondente dovrebbe essere None . Ad esempio, per un'operazione che prende un tensore in virgola mobile x e un indice intero i , la funzione gradiente return [x_grad, None] .

  • Se non esiste alcun gradiente significativo per l'operazione, spesso non sarà necessario registrare alcun gradiente e finché il gradiente dell'operazione non è mai necessario, andrà tutto bene. In alcuni casi, un'operazione non ha un gradiente ben definito ma può essere coinvolta nel calcolo del gradiente. Qui puoi usare ops.NotDifferentiable per propagare automaticamente gli zeri all'indietro.

Si noti che nel momento in cui viene chiamata la funzione gradiente, è disponibile solo il grafico del flusso di dati di ops, non i dati del tensore stesso. Pertanto, tutti i calcoli devono essere eseguiti utilizzando altre operazioni tensorflow, da eseguire al momento dell'esecuzione del grafico.

Aggiungi suggerimenti sul tipo quando registri il gradiente personalizzato per un tipo di operazione per rendere il codice più leggibile, debuggabile, più facile da mantenere e più robusto attraverso la convalida dei dati. Ad esempio, quando prendi un op come parametro in una funzione, specifica che la funzione gradiente prenderà tf.Operation come tipo di parametro.

Funzioni di forma in C++

L'API TensorFlow dispone di una funzionalità chiamata "inferenza della forma" che fornisce informazioni sulle forme dei tensori senza dover eseguire il grafico. L'inferenza della forma è supportata da "funzioni di forma" registrate per ogni tipo di operazione nella dichiarazione REGISTER_OP di C++ e che svolgono due ruoli: affermare che le forme degli input sono compatibili durante la costruzione del grafico e specificare le forme per gli output.

Le funzioni di forma sono definite come operazioni sulla classe shape_inference::InferenceContext . Ad esempio, nella funzione di forma per ZeroOut:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); dichiara che la forma del primo output deve essere impostata sulla forma del primo input. Se l'output è selezionato dal suo indice come nell'esempio precedente, il secondo parametro di set_output dovrebbe essere un oggetto ShapeHandle . Puoi creare un oggetto ShapeHandle vuoto tramite il suo costruttore predefinito. L'oggetto ShapeHandle per un input con indice idx può essere ottenuto da c->input(idx) .

Esistono numerose funzioni di forma comuni che si applicano a molte operazioni, come shape_inference::UnchangedShape che può essere trovata in common_shape_fns.h e utilizzata come segue:

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

Una funzione di forma può anche vincolare la forma di un input. Per la versione di ZeroOut con vincolo di forma vettoriale , la funzione di forma sarebbe la seguente:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));
      c->set_output(0, input);
      return Status::OK();
    });

La chiamata WithRank verifica che la forma di input c->input(0) abbia una forma con esattamente una dimensione (o se la forma di input è sconosciuta, la forma di output sarà un vettore con una dimensione sconosciuta).

Se la tua operazione è polimorfica con più input , puoi utilizzare i membri di InferenceContext per determinare il numero di forme da controllare e Merge per verificare che le forme siano tutte compatibili (in alternativa, accedi agli attributi che indicano le lunghezze, con InferenceContext::GetAttr , che fornisce l'accesso agli attributi dell'op).

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      ::tensorflow::shape_inference::ShapeHandle output;
      for (size_t i = 0; i < c->num_inputs(); ++i) {
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));
        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));
      }
      c->set_output(0, output);
      return Status::OK();
    });

Poiché l'inferenza della forma è una caratteristica opzionale e le forme dei tensori possono variare dinamicamente, le funzioni di forma devono essere robuste per informazioni sulla forma incomplete per qualsiasi input. Il metodo Merge in InferenceContext consente al chiamante di affermare che due forme sono identiche, anche se una o entrambe non dispongono di informazioni complete. Le funzioni di forma sono definite per tutte le operazioni principali di TensorFlow e forniscono molti esempi di utilizzo diversi.

La classe InferenceContext dispone di una serie di funzioni che possono essere utilizzate per definire le manipolazioni delle funzioni di forma. Ad esempio, puoi verificare che una dimensione particolare abbia un valore molto specifico utilizzando InferenceContext::Dim e InferenceContext::WithValue ; puoi specificare che una dimensione di output è la somma/prodotto di due dimensioni di input utilizzando InferenceContext::Add e InferenceContext::Multiply . Consulta la classe InferenceContext per tutte le varie manipolazioni delle forme che puoi specificare. L'esempio seguente imposta la forma del primo output su (n, 3), dove il primo input ha forma (n, ...)

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

Se disponi di una funzione di forma complicata, dovresti prendere in considerazione l'aggiunta di un test per verificare che varie combinazioni di forme di input producano le combinazioni di forme di output previste. Puoi vedere esempi di come scrivere questi test in alcuni dei nostri test operativi principali . (La sintassi di INFER_OK e INFER_ERROR è un po' criptica, ma cerca di essere compatto nel rappresentare le specifiche della forma di input e output nei test. Per ora, vedi i commenti circostanti in quei test per avere un'idea della specifica della stringa di forma).

Crea un pacchetto pip per la tua operazione personalizzata

Per creare un pacchetto pip per la tua operazione, consulta l'esempio tensorflow/custom-op . Questa guida mostra come creare operazioni personalizzate dal pacchetto pip TensorFlow invece di creare TensorFlow dal sorgente.