Crea un'operazione

Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

Se desideri creare un'operazione che non è 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'op C++ personalizzata. Ci sono diversi motivi per cui potresti voler creare un'operazione C++ personalizzata:

  • Non è facile o possibile esprimere la tua 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 avrebbe difficoltà a fondere.

Ad esempio, immagina di voler implementare qualcosa come "median pooling", simile all'operatore "MaxPool", ma calcolando le mediane su finestre scorrevoli anziché sui valori massimi. Potrebbe essere possibile eseguire questa operazione 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 innanzitutto la pena provare ad 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 il nuovo op 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 anche 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 che hai registrato nel passaggio 1. Possono esserci più kernel per diversi tipi o architetture di input / output (ad esempio CPU, GPU).
  3. Crea un wrapper Python (opzionale). Questo wrapper è l'API pubblica utilizzata per creare l'operazione in Python. Un wrapper predefinito viene generato dalla registrazione dell'operazione, che può essere utilizzato direttamente o aggiunto.
  4. Scrivi una funzione per calcolare i gradienti per l'op (opzionale).
  5. Prova l'op. 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 verifica le funzioni forward di operatori simili a Relu e i loro gradienti.

Prerequisiti

Definire l'interfaccia operativa

Definisci l'interfaccia di un'operazione registrandola con il sistema TensorFlow. Nella registrazione, specifichi il nome del tuo op, i suoi input (tipi e nomi) e output (tipi e nomi), nonché docstrings e qualsiasi attr che l'op potrebbe richiedere.

Per vedere come funziona, supponiamo di voler creare un op che accetta un tensore di int32 s e restituisce una copia del tensore, con tutto tranne il primo elemento impostato su zero. Per fare ciò, crea un file chiamato zero_out.cc . Quindi aggiungi una chiamata alla macro REGISTER_OP che definisce l'interfaccia per il tuo op:

#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 ZeroOut prende un tensore to_zero di interi a 32 bit come input e restituisce un tensore zeroed di interi a 32 bit. L'op 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], 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 di context di tipo OpKernelContext* , da cui è possibile accedere a elementi utili come i tensori di input e di 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, specifichi diversi vincoli in base ai quali verrà eseguito questo kernel. Ad esempio, potresti avere un kernel creato per le CPU e uno separato per le GPU.

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

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

Kernel di 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 utilizzati 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 avvio.

A volte l'implementazione di OpKernel è comune tra una CPU e un kernel GPU, ad esempio durante l'ispezione degli input e l'allocazione degli output. In tal caso, un'implementazione suggerita consiste nel:

  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 di functor basata su modelli.
  3. La specializzazione di quel functor 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'op usando il tuo 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 di cui hai bisogno 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 usare per compilare il tuo op 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

In macOS, è necessario il flag aggiuntivo "-undefined dynamic_lookup" durante la creazione del file .so .

Nota sulla versione gcc >=5 : gcc usa il nuovo C++ ABI dalla versione 5 . I pacchetti binari pip disponibili sul sito Web TensorFlow sono creati con gcc4 che utilizza l'ABI precedente. Se compili la tua libreria op con gcc>=5 , aggiungi -D_GLIBCXX_USE_CXX11_ABI=0 alla riga di comando per rendere la libreria compatibile con l'abi precedente.

Compila l'operazione utilizzando bazel (installazione sorgente TensorFlow)

Se hai installato i sorgenti TensorFlow, puoi utilizzare il sistema di build di TensorFlow per compilare il tuo op. Inserisci un file BUILD con la seguente regola di build Bazel nella 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 di Example , con il kernel CUDA, è necessario utilizzare il parametro gpu_srcs di tf_custom_op_library . Posiziona un file BUILD con la seguente regola di build Bazel in una nuova cartella all'interno della 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 kernel_example.so .

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

Usa l'op in Python

L'API Python di TensorFlow 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'op e il kernel. Pertanto, una volta creato l'op, puoi eseguire le seguenti operazioni per eseguirlo 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 ). Quindi, se il tuo op è chiamato ZeroOut nei file C++, la funzione python sarà chiamata zero_out .

Per rendere disponibile l'op 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

Verifica che l'operazione funzioni

Un buon modo per verificare di aver implementato correttamente la tua operazione è scrivere un test per esso. 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 tensorflow sia installato):

$ 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 limitate), esamineremo alcune delle cose più complicate che in genere dovrai integrare nella tua operazione. Ciò comprende:

Controlli condizionali e validazione

L'esempio sopra presupponeva che l'op si applicasse a un tensore di qualsiasi forma. E se si applicasse solo ai vettori? Ciò significa aggiungere un controllo all'implementazione di OpKernel 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."));
    // ...
  }

Questo afferma che l'input è un vettore e restituisce dopo aver impostato lo stato InvalidArgument se non lo è. La macro OP_REQUIRES accetta tre argomenti:

In alternativa, se si desidera verificare se un oggetto Status restituito da una funzione è un errore e, in tal caso, restituirlo, utilizzare OP_REQUIRES_OK . Entrambe queste macro ritornano dalla funzione in caso di errore.

Registrazione operativa

Attr

Le operazioni possono avere attrs, i cui valori vengono impostati quando l'operazione viene aggiunta a un grafico. Questi sono usati per configurare l'op e ai loro valori è possibile accedere sia all'interno dell'implementazione del kernel che nei tipi di input e output nella registrazione dell'op. Preferisci usare un input invece di 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 usando un feed, ecc. Gli Attr sono usati per cose che non possono essere fatte con gli input: qualsiasi configurazione che influisca sulla firma (numero o tipo di input o output) o che puo' t cambiare da un passo all'altro.

Definisci un attr quando registri l'op, specificandone il nome e il tipo usando 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 del modulo descritto di seguito .

Ad esempio, se desideri che l' 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 attributo è diverso dal tf.DType utilizzato per gli input e gli output.)

Il tuo kernel può quindi accedere a questo attr nel suo costruttore tramite il parametro di 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 : un TensorShapeProto .
  • list(<type>) : un elenco di <type> , dove <type> è uno dei tipi precedenti. Si noti che list(list(<type>)) non è valido.

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

Valori e vincoli predefiniti

Attrs può avere valori predefiniti e alcuni tipi di attrs possono avere vincoli. Per definire un attr con vincoli, è possibile utilizzare i seguenti <attr-type-expr> s:

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

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 . Questo è implicito quando hai un elenco di tipi in {...} . Ad esempio, in questo caso attr t è un tipo che deve essere an int32 , a float o a bool :

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

Esistono scorciatoie per i vincoli di tipo comuni:

  • numbertype : tipo di type 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'attr 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

Le liste possono essere combinate con altre liste e singole tipologie. Il seguente op 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 l'attr a deve avere un valore che sia 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 opzionale 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 proto della definizione GraphDef risultante.

Di seguito sono riportati 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

Tipo polimorfismo

Per le operazioni che possono accettare tipi diversi come input o produrre tipi di output diversi, è possibile specificare un attr in un tipo di input o output nella registrazione dell'operazione. In genere si registra quindi un OpKernel per ogni tipo supportato.

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

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

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

Denominazione

In genere, a input, output e attr devono essere assegnati nomi snake_case. L'unica eccezione è attrs che viene utilizzato come tipo di input o nel tipo di output. Tali attr possono essere dedotti quando l'op viene aggiunto al grafico e quindi non compaiono nella funzione dell'op. 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 ). Agli attributi dedotti 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`.
  """
Digitare esempio di polimorfismo
#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 op esistente:

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

Supponiamo che tu voglia 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 un modello C++. Avrai ancora 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 overload, 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 di 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 ad essere in grado di accettare o produrre tipi diversi, le operazioni 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 dell'input in che dell'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");

Puoi anche porre restrizioni sui tipi che possono essere specificati nell'elenco. Nel prossimo caso, l'input è un elenco di float e double tensors. L'op 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.

Questo può essere reso anche di tipo 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 di tensori hanno una lunghezza minima di 1. È possibile modificare tale impostazione predefinita utilizzando un vincolo ">=" sul corrispondente attr . In questo 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 "list(type)" attrs:

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

Ingressi e uscite

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

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

Ogni 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 tipo type o list(type) (con una possibile restrizione di 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 attr 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 ei tipi di tensori nell'output out sono gli stessi dell'input in , poiché entrambi sono di tipo T

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

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

    Considerando che questa op 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 attr utilizzato nel tipo di input. Per convenzione quegli attr dedotti usano nomi in maiuscolo (come T o N ). Altrimenti input, output e attrs hanno nomi come parametri di funzione (ad esempio num_outputs ). Per maggiori dettagli, vedere la sezione precedente sulla denominazione .

Per maggiori 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 felici che usano la tua operazione. Tuttavia, vorresti apportare modifiche all'operazione in qualche modo.

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

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

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

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

    puoi renderlo polimorfico in un modo compatibile con le versioni precedenti usando:

    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: T")
        .Output("out: T")
        .Attr("T: numerictype = DT_FLOAT");
    
  2. Puoi tranquillamente fare un vincolo su un attr meno restrittivo. Ad esempio, puoi passare 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 ingressi/uscite in ingressi/uscite elenco, purché l'impostazione predefinita per il tipo di elenco corrisponda alla vecchia firma.

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

  5. Namespace 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 qualsiasi operazione che potrebbe essere inclusa nelle versioni future di TensorFlow.

  6. Pianificare in anticipo! Cerca di anticipare usi futuri dell'op. Alcune modifiche alla firma non possono essere eseguite in modo compatibile (ad esempio, trasformando 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 apportare la modifica a un'operazione compatibile con le versioni precedenti, creare una nuova operazione con un nuovo nome con la nuova semantica.

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

Supporto GPU

Puoi implementare diversi OpKernel e registrarne uno per CPU e un altro per GPU, proprio come puoi registrare kernel per tipi diversi . Esistono diversi esempi di kernel con supporto GPU in tensorflow/core/kernels/ . Si noti 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 del codice condiviso in comune in un file .h .

Ad esempio, tf.pad ha tutto tranne il kernel della GPU in tensorflow/core/kernels/pad_op.cc . Il kernel della 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 di pad , ha ancora bisogno del suo input "paddings" nella memoria della CPU. Per contrassegnare che gli input o gli output sono mantenuti sulla CPU, aggiungere 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 op. La 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'uso 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 usare 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, usando la funzione tf.load_op_library .

Nota 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 il tuo 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 fare in modo che la differenziazione automatica funzioni per le nuove operazioni, è necessario registrare una funzione gradiente che calcola le sfumature rispetto agli input delle operazioni dati i gradienti rispetto agli output delle operazioni.

Matematicamente, se un op calcola \(y = f(x)\) il gradiente registrato op converte i gradienti \(\partial L/ \partial y\) della perdita \(L\) rispetto a\(y\) in gradienti \(\partial L/ \partial x\) rispetto 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". Questo è 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 costruirà nuove operazioni dai tensori op.inputs[i] , op.outputs[i] e grad . Le informazioni su qualsiasi attr possono essere trovate tramite tf.Operation.get_attr .

  • Se l'op 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, ad esempio per input interi usati 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 c'è alcun gradiente significativo per l'operazione, spesso non dovrai registrare alcun gradiente e finché il gradiente dell'operazione non è mai necessario, andrà tutto bene. In alcuni casi, un op non ha un gradiente ben definito ma può essere coinvolto 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 delle operazioni, non i dati del tensore stesso. Pertanto, tutto il calcolo deve essere eseguito utilizzando altre operazioni di flusso tensoriale, da eseguire al momento dell'esecuzione del grafico.

Funzioni di forma in C++

L'API TensorFlow ha una funzione chiamata "shape inference" che fornisce informazioni sulle forme dei tensori senza dover eseguire il grafico. L'inferenza di forma è supportata da "funzioni di forma" registrate per ogni tipo op nella dichiarazione C++ REGISTER_OP e svolgono due ruoli: affermando che le forme degli input sono compatibili durante la costruzione del grafico e specificando le forme per gli output.

Le funzioni Shape sono definite come operazioni sulla classe shape_inference::InferenceContext . Ad esempio, nella funzione 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 in base al suo costruttore predefinito. L'oggetto ShapeHandle per un input con index 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 trovato in common_shape_fns.h e utilizzato 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 un 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 convalida 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 il tuo op è polimorfico con più input , puoi usare i membri di InferenceContext per determinare il numero di forme da controllare e Merge per convalidare 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 di forma è una funzionalità 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 uguali, anche se una o entrambe non dispongono di informazioni complete. Le funzioni Shape sono definite per tutte le operazioni principali di TensorFlow e forniscono molti esempi di utilizzo diversi.

La classe InferenceContext dispone di numerose funzioni che possono essere utilizzate per definire le manipolazioni delle funzioni di forma. Ad esempio, puoi verificare che una determinata dimensione abbia un valore molto specifico utilizzando InferenceContext::Dim e InferenceContext::WithValue ; è possibile specificare che una dimensione di output è la somma/prodotto di due dimensioni di input utilizzando InferenceContext::Add e InferenceContext::Multiply . Vedi la classe InferenceContext per tutte le varie manipolazioni di forme che puoi specificare. L'esempio seguente imposta la forma del primo output su (n, 3), dove il primo input ha la 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 hai una funzione di forma complicata, dovresti considerare di aggiungere un test per convalidare che varie combinazioni di forme di input producono 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, guarda 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, vedi l' esempio tensorflow/custom-op . Questa guida mostra come creare operazioni personalizzate dal pacchetto pip TensorFlow invece di creare TensorFlow dal sorgente.