Esta página se ha traducido con Cloud Translation API.

crear una operación

Si desea crear una operación que no esté cubierta por la biblioteca TensorFlow existente, le recomendamos que primero intente escribir la operación en Python como una composición de operaciones o funciones de Python existentes. Si eso no es posible, puede crear una operación C++ personalizada. Hay varias razones por las que es posible que desee crear una operación C++ personalizada:

No es fácil ni posible expresar su operación como una composición de operaciones existentes.
No es eficiente expresar su operación como una composición de primitivas existentes.
Quiere fusionar manualmente una composición de primitivas que a un futuro compilador le resultaría difícil fusionar.

Por ejemplo, imagine que desea implementar algo como "agrupación de medianas", similar al operador "MaxPool", pero calculando medianas sobre ventanas deslizantes en lugar de valores máximos. Hacer esto usando una composición de operaciones puede ser posible (por ejemplo, usando ExtractImagePatches y TopK), pero puede no ser tan eficiente en rendimiento o memoria como una operación nativa donde se puede hacer algo más inteligente en una única operación fusionada. Como siempre, normalmente primero vale la pena intentar expresar lo que desea utilizando la composición de operadores, y solo elegir agregar una nueva operación si resulta difícil o ineficiente.

Para incorporar su operación personalizada necesitará:

Registre la nueva operación en un archivo C++. El registro de operación define una interfaz (especificación) para la funcionalidad de la operación, que es independiente de la implementación de la operación. Por ejemplo, el registro de operación define el nombre de la operación y las entradas y salidas de la operación. También define la función de forma que se utiliza para la inferencia de formas tensoriales.
Implemente la operación en C++. La implementación de una operación se conoce como kernel y es la implementación concreta de la especificación que registró en el Paso 1. Puede haber múltiples kernels para diferentes tipos o arquitecturas de entrada/salida (por ejemplo, CPU, GPU).
Cree un contenedor de Python (opcional). Este contenedor es la API pública que se utiliza para crear la operación en Python. Se genera un contenedor predeterminado a partir del registro de operación, que se puede usar directamente o agregar.
Escriba una función para calcular gradientes para la operación (opcional).
Pruebe la operación. Normalmente hacemos esto en Python por conveniencia, pero también puedes probar la operación en C++. Si define gradientes, puede verificarlos con Python tf.test.compute_gradient_error . Vea relu_op_test.py como un ejemplo que prueba las funciones directas de los operadores similares a Relu y sus gradientes.

Requisitos previos

Algo de familiaridad con C++.
Debe haber instalado el binario de TensorFlow o debe haber descargado el código fuente de TensorFlow y poder compilarlo.

Definir la interfaz operativa.

La interfaz de una operación se define registrándola en el sistema TensorFlow. En el registro, usted especifica el nombre de su operación, sus entradas (tipos y nombres) y salidas (tipos y nombres), así como cadenas de documentos y cualquier atributo que la operación pueda requerir.

Para ver cómo funciona esto, suponga que desea crear una operación que tome un tensor de int32 sy genere una copia del tensor, con todos los elementos excepto el primero establecidos en cero. Para hacer esto, cree un archivo llamado zero_out.cc . Luego agregue una llamada a la macro REGISTER_OP que define la interfaz para su operación:

#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();
    });

Esta operación ZeroOut toma un tensor to_zero de enteros de 32 bits como entrada y genera un tensor zeroed de enteros de 32 bits. La operación también utiliza una función de forma para garantizar que el tensor de salida tenga la misma forma que el tensor de entrada. Por ejemplo, si la entrada es un tensor de forma [10, 20], entonces esta función de forma especifica que la forma de salida también es [10, 20].

Implementar el kernel para la operación.

Después de definir la interfaz, proporcione una o más implementaciones de la operación. Para crear uno de estos núcleos, cree una clase que extienda OpKernel y anule el método Compute . El método Compute proporciona un argumento context de tipo OpKernelContext* , desde el cual puede acceder a cosas útiles como los tensores de entrada y salida.

Agregue su kernel al archivo que creó anteriormente. El kernel podría verse así:

#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);
  }
};

Después de implementar su kernel, lo registra en el sistema TensorFlow. En el registro, usted especifica diferentes restricciones bajo las cuales se ejecutará este kernel. Por ejemplo, es posible que tenga un núcleo diseñado para CPU y otro separado para GPU.

Para hacer esto para la operación ZeroOut , agregue lo siguiente a zero_out.cc :

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

Importante: Se puede acceder a instancias de su OpKernel simultáneamente. Su método Compute debe ser seguro para subprocesos. Proteja cualquier acceso a los miembros de la clase con un mutex. O mejor aún, ¡no comparta el estado a través de los miembros de la clase! Considere utilizar ResourceMgr para realizar un seguimiento del estado de operación.

Núcleos de CPU multiproceso

Para escribir un núcleo de CPU multiproceso, se puede utilizar la función Shard en work_sharder.h . Esta función fragmenta una función de cálculo entre los subprocesos configurados para usarse para subprocesos intraoperativos (consulte intra_op_parallelism_threads en config.proto ).

Núcleos de GPU

Un kernel de GPU se implementa en dos partes: el OpKernel y el kernel CUDA y su código de lanzamiento.

A veces, la implementación de OpKernel es común entre un kernel de CPU y GPU, como cuando se trata de inspeccionar entradas y asignar salidas. En ese caso, una implementación sugerida es:

Defina la plantilla OpKernel en el Dispositivo y el tipo primitivo del tensor.
Para realizar el cálculo real de la salida, la función Compute llama a una estructura funtor con plantilla.
La especialización de ese funtor para CPUDevice se define en el mismo archivo, pero la especialización para GPUDevice se define en un archivo .cu.cc, ya que se compilará con el compilador CUDA.

A continuación se muestra una implementación de ejemplo.

// 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

Construir la biblioteca de operaciones

Compile la operación usando el compilador de su sistema (instalación binaria de TensorFlow)

Debería poder compilar zero_out.cc con un compilador de C++ como g++ o clang disponible en su sistema. El paquete PIP binario instala los archivos de encabezado y la biblioteca que necesita para compilar su operación en ubicaciones específicas del sistema. Sin embargo, la biblioteca Python de TensorFlow proporciona la función get_include para obtener el directorio del encabezado, y el directorio get_lib tiene un objeto compartido para vincular. Aquí están los resultados de estas funciones en una máquina 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'

Suponiendo que tiene g++ instalado, aquí está la secuencia de comandos que puede usar para compilar su operación en una biblioteca dinámica.

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

En macOS, se requiere la marca adicional "-undefineddynamic_lookup" al crear el archivo .so .

Nota sobre la versión gcc >=5 : gcc usa la nueva ABI de C++ desde la versión 5 . TensorFlow 2.8 y versiones anteriores se crearon con gcc4 que utiliza la ABI anterior. Si está utilizando estas versiones de TensorFlow y está intentando compilar su biblioteca de operaciones con gcc>=5 , agregue -D_GLIBCXX_USE_CXX11_ABI=0 a la línea de comando para que la biblioteca sea compatible con la ABI anterior. Los paquetes de TensorFlow 2.9+ son compatibles con la ABI más nueva de forma predeterminada.

Compile la operación usando bazel (instalación fuente de TensorFlow)

Si tiene fuentes de TensorFlow instaladas, puede utilizar el sistema de compilación de TensorFlow para compilar su operación. Coloque un archivo BUILD con la siguiente regla de compilación de Bazel en el directorio tensorflow/core/user_ops .

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

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

Ejecute el siguiente comando para compilar zero_out.so .

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

Para compilar la operación Example , con el kernel CUDA, debe usar el parámetro gpu_srcs de tf_custom_op_library . Coloque un archivo BUILD con la siguiente regla de compilación de Bazel en una nueva carpeta dentro del directorio tensorflow/core/user_ops (por ejemplo, "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"],
)

Ejecute el siguiente comando para compilar kernel_example.so .

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

Nota: Aunque puede crear una biblioteca compartida (un archivo .so ) con la regla cc_library estándar, le recomendamos encarecidamente que utilice la macro tf_custom_op_library . Agrega algunas dependencias requeridas y realiza comprobaciones para garantizar que la biblioteca compartida sea compatible con el mecanismo de carga de complementos de TensorFlow.

Usa la operación en Python

La API de TensorFlow Python proporciona la función tf.load_op_library para cargar la biblioteca dinámica y registrar la operación con el marco de TensorFlow. load_op_library devuelve un módulo de Python que contiene los contenedores de Python para la operación y el kernel. Por lo tanto, una vez que haya creado la operación, puede hacer lo siguiente para ejecutarla desde 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)

Tenga en cuenta que a la función generada se le asignará un nombre Snake_case (para cumplir con PEP8 ). Entonces, si su operación se llama ZeroOut en los archivos C++, la función de Python se llamará zero_out .

Para que la operación esté disponible como una función normal que se pueda import desde un módulo de Python, puede ser útil tener la llamada load_op_library en un archivo fuente de Python de la siguiente manera:

import tensorflow as tf

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

Verifique que la operación funcione

Una buena manera de verificar que ha implementado exitosamente su operación es escribir una prueba para ella. Cree el archivo zero_out_op_test.py con el contenido:

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()

Luego ejecute su prueba (suponiendo que tenga tensorflow instalado):

$ python zero_out_op_test.py

Incorpore funciones avanzadas a su operación

Ahora que sabe cómo crear una operación e implementación básica (y algo restringida), veremos algunas de las cosas más complicadas que normalmente necesitará incorporar a su operación. Esto incluye:

Verificaciones condicionales y validación.
Registro de operaciones
Soporte de GPU
- Compilando el kernel para el dispositivo GPU
Implementar el gradiente en Python
Funciones de forma en C++

Verificaciones condicionales y validación.

El ejemplo anterior asumió que la operación se aplicaba a un tensor de cualquier forma. ¿Qué pasaría si solo se aplicara a vectores? Eso significa agregar una verificación a la implementación de OpKernel anterior.

  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."));
    // ...
  }

Esto afirma que la entrada es un vector y regresa habiendo establecido el estado InvalidArgument si no lo es. La macro OP_REQUIRES toma tres argumentos:

El context , que puede ser un puntero OpKernelContext u OpKernelConstruction (consulte tensorflow/core/framework/op_kernel.h ), para su método SetStatus() .
La condición. Por ejemplo, hay funciones para validar la forma de un tensor en tensorflow/core/framework/tensor_shape.h
El error en sí, que está representado por un objeto Status , consulte tensorflow/core/platform/status.h . Un Status tiene un tipo (frecuentemente InvalidArgument , pero consulte la lista de tipos) y un mensaje. Las funciones para construir un error se pueden encontrar en tensorflow/core/platform/errors.h .

Alternativamente, si desea probar si un objeto Status devuelto por alguna función es un error y, de ser así, devolverlo, use OP_REQUIRES_OK . Ambas macros regresan de la función en caso de error.

Registro de operaciones

atributos

Las operaciones pueden tener atributos, cuyos valores se establecen cuando la operación se agrega a un gráfico. Estos se utilizan para configurar la operación, y se puede acceder a sus valores tanto dentro de la implementación del kernel como en los tipos de entradas y salidas en el registro de la operación. Prefiera usar una entrada en lugar de un atributo cuando sea posible, ya que las entradas son más flexibles. Esto se debe a que los atributos son constantes y deben definirse en el momento de la construcción del gráfico. Por el contrario, las entradas son tensores cuyos valores pueden ser dinámicos; es decir, las entradas pueden cambiar en cada paso, configurarse mediante un feed, etc. Los atributos se usan para cosas que no se pueden hacer con las entradas: cualquier configuración que afecte la firma (número o tipo de entradas o salidas) o que pueda t cambiar de un paso a otro.

Usted define un atributo cuando registra la operación, especificando su nombre y tipo usando el método Attr , que espera una especificación del formulario:

<name>: <attr-type-expr>

donde <name> comienza con una letra y puede estar compuesto por caracteres alfanuméricos y guiones bajos, y <attr-type-expr> es una expresión de tipo del formato que se describe a continuación .

Por ejemplo, si desea que la operación ZeroOut conserve un índice especificado por el usuario, en lugar de solo el elemento 0, puede registrar la operación de esta manera:

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

(Tenga en cuenta que el conjunto de tipos de atributos es diferente del tf.DType utilizado para entradas y salidas).

Luego, su kernel puede acceder a este atributo en su constructor a través del parámetro 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_;
};

que luego se puede utilizar en el método 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_);
  }

Tipos de atributos

Los siguientes tipos son compatibles con un atributo:

string : cualquier secuencia de bytes (no es necesario que sea UTF8).
int : un entero con signo.
float : un número de punto flotante.
bool : verdadero o falso.
type : uno de los valores (sin referencia) de DataType .
shape : Un TensorShapeProto .
list(<type>) : una lista de <type> , donde <type> es uno de los tipos anteriores. Tenga en cuenta que list(list(<type>)) no es válido.

Consulte también: op_def_builder.cc:FinalizeAttr para obtener una lista definitiva.

Valores predeterminados y restricciones

Los atributos pueden tener valores predeterminados y algunos tipos de atributos pueden tener restricciones. Para definir un atributo con restricciones, puede utilizar los siguientes <attr-type-expr> s:

{'<string1>', '<string2>'} : el valor debe ser una cadena que tenga el valor <string1> o <string2> . El nombre del tipo, string , está implícito cuando se utiliza esta sintaxis. Esto emula una enumeración:

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

{<type1>, <type2>} : el valor es de tipo type y debe ser uno de <type1> o <type2> , donde <type1> y <type2> son compatibles con tf.DType . No especifica que el tipo de atributo sea type . Esto está implícito cuando tiene una lista de tipos en {...} . Por ejemplo, en este caso el t es un tipo que debe ser int32 , float o bool :

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

Existen atajos para restricciones de tipos comunes:

numbertype : type restringido a los tipos numéricos (que no son cadenas ni booleanos).
realnumbertype : como numbertype sin tipos complejos.
quantizedtype : como numbertype pero solo los tipos de números cuantificados.

Las listas específicas de tipos permitidos por estos están definidas por las funciones (como NumberTypes() ) en tensorflow/core/framework/types.h . En este ejemplo, el t debe ser uno de los tipos numéricos:

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

Para esta operación:

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

Las listas se pueden combinar con otras listas y tipos individuales. La siguiente operación permite que attr t sea cualquiera de los tipos numéricos o del tipo bool:

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

Para esta operación:

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> : El valor debe ser un int cuyo valor sea mayor o igual a <n> , donde <n> es un número natural. Por ejemplo, el siguiente registro de operación especifica que el atributo a tener un valor que sea al menos 2 :

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

list(<type>) >= <n> : una lista de tipo <type> cuya longitud es mayor o igual a <n> . Por ejemplo, el siguiente registro de operación especifica que el atributo a es una lista de tipos (ya sea int32 o float ) y que debe haber al menos 3 de ellos:

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

Para establecer un valor predeterminado para un atributo (haciéndolo opcional en el código generado), agregue = <default> al final, como en:

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

Además, se pueden especificar tanto una restricción como un valor predeterminado:

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

La sintaxis admitida del valor predeterminado es la que se usaría en la representación proto de la definición GraphDef resultante.

A continuación se muestran ejemplos de cómo especificar un valor predeterminado para todos los tipos:

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]");

Tenga en cuenta en particular que los valores de tipo type utilizan tf.DType .

Polimorfismo

Tipo polimorfismo

Para operaciones que pueden tomar diferentes tipos como entrada o producir diferentes tipos de salida, puede especificar un atributo en un tipo de entrada o salida en el registro de operación. Normalmente, luego registraría un OpKernel para cada tipo admitido.

Por ejemplo, si desea que la operación ZeroOut funcione en float s además de int32 s, el registro de su operación podría verse así:

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

Su registro de operación ahora especifica que el tipo de entrada debe ser float o int32 y que su salida será del mismo tipo, ya que ambos tienen tipo T

Nombrar

Las entradas, salidas y atributos generalmente deben recibir nombres de casos de serpiente. La única excepción son los atributos que se utilizan como tipo de entrada o como tipo de salida. Esos atributos se pueden inferir cuando la operación se agrega al gráfico y, por lo tanto, no aparecen en la función de la operación. Por ejemplo, esta última definición de ZeroOut generará una función de Python similar a la siguiente:

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`.
  """

Si to_zero se le pasa un tensor int32 , entonces T se establece automáticamente en int32 (bueno, en realidad DT_INT32 ). Esos atributos inferidos reciben nombres en mayúscula o CamelCase.

Compare esto con una operación que tiene un atributo de tipo que determina el tipo de salida:

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");

En este caso, el usuario debe especificar el tipo de salida, como en el Python generado:

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`.
  """

Ejemplo de polimorfismo de 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);

Para preservar la compatibilidad con versiones anteriores , debe especificar un valor predeterminado al agregar un atributo a una operación existente:

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

Digamos que desea agregar más tipos, digamos double :

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

En lugar de escribir otro OpKernel con código redundante como el anterior, a menudo podrá utilizar una plantilla C++. Aún tendrá un registro del kernel (llamada REGISTER_KERNEL_BUILDER ) por sobrecarga.

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>);

Si tiene más de un par de sobrecargas, puede colocar el registro en 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

Dependiendo de la lista de tipos para los que esté registrando el kernel, es posible que pueda utilizar una macro proporcionada por 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

Listar entradas y salidas

Además de poder aceptar o producir diferentes tipos, las operaciones pueden consumir o producir un número variable de tensores.

En el siguiente ejemplo, el atributo T contiene una lista de tipos y se utiliza como tipo tanto de in como de out . La entrada y la salida son listas de tensores de ese tipo (y el número y los tipos de tensores en la salida son los mismos que los de la entrada, ya que ambos tienen tipo T ).

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

También puede imponer restricciones sobre los tipos que se pueden especificar en la lista. En el siguiente caso, la entrada es una lista de tensores double y float . La operación acepta, por ejemplo, tipos de entrada (float, double, float) y en ese caso el tipo de salida también sería (float, double, float) .

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

Si desea que todos los tensores de una lista sean del mismo tipo, puede hacer algo como:

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

Esto acepta una lista de tensores int32 y usa un atributo int N para especificar la longitud de la lista.

Esto también se puede hacer tipo polimórfico . En el siguiente ejemplo, la entrada es una lista de tensores (con longitud "N" ) del mismo tipo (pero sin especificar) ( "T" ), y la salida es un único tensor de tipo coincidente:

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

De forma predeterminada, las listas de tensores tienen una longitud mínima de 1. Puede cambiar ese valor predeterminado usando una restricción ">=" en el atributo correspondiente . En el siguiente ejemplo, la entrada es una lista de al menos 2 tensores int32 :

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

La misma sintaxis funciona con los atributos "list(type)" :

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

Entradas y salidas

Para resumir lo anterior, un registro de operación puede tener múltiples entradas y salidas:

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

Cada especificación de entrada o salida tiene la forma:

<name>: <io-type-expr>

donde <name> comienza con una letra y puede estar compuesto por caracteres alfanuméricos y guiones bajos. <io-type-expr> es una de las siguientes expresiones de tipo:

<type> , donde <type> es un tipo de entrada admitido (por ejemplo, float , int32 , string ). Esto especifica un único tensor del tipo dado.
Consulte tf.DType .
```
REGISTER_OP("BuiltInTypesExample")
    .Input("integers: int32")
    .Input("complex_numbers: complex64");
```

<attr-type> , donde <attr-type> es el nombre de un Attr con type tipo o list(type) (con una posible restricción de tipo). Esta sintaxis permite operaciones polimórficas .

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

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

Hacer referencia a un atributo de tipo list(type) le permite aceptar una secuencia de tensores.

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");

Tenga en cuenta que el número y los tipos de tensores en la salida out son los mismos que en la entrada in , ya que ambos son de tipo T .

Para una secuencia de tensores con el mismo tipo: <number> * <type> , donde <number> es el nombre de un Attr de tipo int . El <type> puede ser tf.DType o el nombre de un atributo con tipo type . Como ejemplo del primero, esta operación acepta una lista de tensores int32 :
```
REGISTER_OP("Int32SequenceExample")
    .Attr("NumTensors: int")
    .Input("in: NumTensors * int32")
```
Mientras que esta operación acepta una lista de tensores de cualquier tipo, siempre que sean todos iguales:
```
REGISTER_OP("SameTypeSequenceExample")
    .Attr("NumTensors: int")
    .Attr("T: type")
    .Input("in: NumTensors * T")
```
Para una referencia a un tensor: Ref(<type>) , donde <type> es uno de los tipos anteriores.

Se inferirá cualquier atributo utilizado en el tipo de entrada. Por convención, esos atributos inferidos utilizan nombres en mayúscula (como T o N ). De lo contrario, las entradas, salidas y atributos tienen nombres como parámetros de función (por ejemplo, num_outputs ). Para obtener más detalles, consulte la sección anterior sobre nombres .

Para obtener más detalles, consulte tensorflow/core/framework/op_def_builder.h .

Compatibilidad con versiones anteriores

Supongamos que ha escrito una operación agradable y personalizada y la ha compartido con otros, de modo que tenga clientes satisfechos utilizando su operación. Sin embargo, le gustaría realizar cambios en la operación de alguna manera.

En general, los cambios en las especificaciones registradas existentes deben ser compatibles con versiones anteriores: cambiar la especificación de una operación no debe romper los búferes de protocolo GraphDef serializados anteriores construidos a partir de especificaciones anteriores. Los detalles de la compatibilidad GraphDef se describen aquí .

Hay varias formas de preservar la compatibilidad con versiones anteriores.

Cualquier atributo nuevo agregado a una operación debe tener valores predeterminados definidos y, con ese valor predeterminado, la operación debe tener el comportamiento original. Para cambiar una operación de no polimórfica a polimórfica, debe asignar un valor predeterminado al nuevo tipo de atributo para conservar la firma original de forma predeterminada. Por ejemplo, si su operación fue:
```
REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: float")
    .Output("out: float");
```
puedes hacerlo polimórfico de una manera compatible con versiones anteriores usando:
```
REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: T")
    .Output("out: T")
    .Attr("T: numerictype = DT_FLOAT");
```
Puede hacer que una restricción de un atributo sea menos restrictiva de forma segura. Por ejemplo, puede cambiar de {int32, int64} a {int32, int64, float} o type . O puedes cambiar de {"apple", "orange"} a {"apple", "banana", "orange"} o string .
Puede cambiar entradas/salidas individuales en entradas/salidas de lista, siempre que el valor predeterminado para el tipo de lista coincida con la firma anterior.
Puede agregar una nueva lista de entrada/salida, si por defecto está vacía.
Espacio de nombres para cualquier operación nueva que cree, anteponiendo a los nombres de las operaciones algo exclusivo de su proyecto. Esto evita que su operación colisione con cualquier operación que pueda incluirse en versiones futuras de TensorFlow.
¡Planifique con anticipación! Intente anticipar usos futuros de la operación. Algunos cambios de firma no se pueden realizar de forma compatible (por ejemplo, convertir una lista del mismo tipo en una lista de distintos tipos).

La lista completa de cambios seguros e inseguros se puede encontrar en tensorflow/core/framework/op_compatibility_test.cc . Si no puede hacer que el cambio a una operación sea compatible con versiones anteriores, cree una nueva operación con un nuevo nombre y la nueva semántica.

También tenga en cuenta que, si bien estos cambios pueden mantener la compatibilidad GraphDef , el código Python generado puede cambiar de una manera que no sea compatible con las personas que llaman anteriormente. La API de Python puede mantenerse compatible mediante cambios cuidadosos en un contenedor de Python escrito a mano, manteniendo la firma anterior excepto posiblemente agregando nuevos argumentos opcionales al final. Generalmente, solo se pueden realizar cambios incompatibles cuando TensorFlow cambia las versiones principales y deben ajustarse a la semántica de la versión GraphDef .

Soporte de GPU

Puede implementar diferentes OpKernels y registrar uno para CPU y otro para GPU, del mismo modo que puede registrar kernels para diferentes tipos . Hay varios ejemplos de kernels con soporte para GPU en tensorflow/core/kernels/ . Observe que algunos núcleos tienen una versión de CPU en un archivo .cc , una versión de GPU en un archivo que termina en _gpu.cu.cc y algo de código compartido en un archivo .h .

Por ejemplo, tf.pad tiene todo menos el kernel de GPU en tensorflow/core/kernels/pad_op.cc . El kernel de GPU está en tensorflow/core/kernels/pad_op_gpu.cu.cc y el código compartido es una clase con plantilla definida en tensorflow/core/kernels/pad_op.h . Organizamos el código de esta manera por dos razones: le permite compartir código común entre las implementaciones de CPU y GPU, y coloca la implementación de GPU en un archivo separado para que solo pueda ser compilado por el compilador de GPU.

Una cosa a tener en cuenta, incluso cuando se utiliza la versión de pad del kernel de GPU, todavía necesita su entrada "paddings" en la memoria de la CPU. Para marcar que las entradas o salidas se mantienen en la CPU, agregue una llamada HostMemory() al registro del kernel, por ejemplo:

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

Compilando el kernel para el dispositivo GPU

Mire cuda_op_kernel.cu.cc para ver un ejemplo que utiliza un kernel CUDA para implementar una operación. tf_custom_op_library acepta un argumento gpu_srcs en el que se puede especificar la lista de archivos fuente que contienen los núcleos CUDA (archivos *.cu.cc ). Para usar con una instalación binaria de TensorFlow, los núcleos CUDA deben compilarse con el compilador nvcc de NVIDIA. Aquí está la secuencia de comandos que puede usar para compilar cuda_op_kernel.cu.cc y cuda_op_kernel.cc en una única biblioteca cargable dinámicamente:

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 producido anteriormente se puede cargar como de costumbre en Python, usando la función tf.load_op_library .

Tenga en cuenta que si sus bibliotecas CUDA no están instaladas en /usr/local/lib64 , deberá especificar la ruta explícitamente en el segundo comando (g++) anterior. Por ejemplo, agregue -L /usr/local/cuda-8.0/lib64/ si su CUDA está instalado en /usr/local/cuda-8.0 .

Implementar el gradiente en Python

Dado un gráfico de operaciones, TensorFlow usa diferenciación automática (propagación hacia atrás) para agregar nuevas operaciones que representen gradientes con respecto a las operaciones existentes. Para que la diferenciación automática funcione para nuevas operaciones, debe registrar una función de gradiente que calcule los gradientes con respecto a las entradas de las operaciones, dados los gradientes con respecto a las salidas de las operaciones.

Matemáticamente, si una operación calcula \(y = f(x)\) la operación de gradiente registrada convierte gradientes \(\partial L/ \partial y\) de perdida \(L\) con respecto a\(y\) en gradientes \(\partial L/ \partial x\) con respecto a \(x\) mediante la regla de la cadena:

\[\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}.\]

En el caso de ZeroOut , solo una entrada en la entrada afecta la salida, por lo que el gradiente con respecto a la entrada es un tensor escaso "one hot". Esto se expresa de la siguiente manera:

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

Detalles sobre el registro de funciones de gradiente con tf.RegisterGradient :

Para una operación con una salida, la función de gradiente tomará un tf.Operation , op y un tf.Tensor grad y construirá nuevas operaciones a partir de los tensores op.inputs[i] , op.outputs[i] y grad . La información sobre cualquier atributo se puede encontrar a través de tf.Operation.get_attr .
Si la operación tiene múltiples salidas, la función de gradiente tomará op y grads , donde grads es una lista de gradientes con respecto a cada salida. El resultado de la función de gradiente debe ser una lista de objetos Tensor que representen los gradientes con respecto a cada entrada.
Si no hay un gradiente bien definido para alguna entrada, como para las entradas enteras utilizadas como índices, el gradiente devuelto correspondiente debe ser None . Por ejemplo, para una operación que toma un tensor de punto flotante x y un índice entero i , la función de gradiente return [x_grad, None] .
Si no hay ningún gradiente significativo para la operación, a menudo no tendrá que registrar ningún gradiente y, siempre que el gradiente de la operación nunca sea necesario, estará bien. En algunos casos, una operación no tiene un gradiente bien definido pero puede participar en el cálculo del gradiente. Aquí puedes usar ops.NotDifferentiable para propagar automáticamente los ceros hacia atrás.

Tenga en cuenta que en el momento en que se llama a la función de gradiente, solo está disponible el gráfico de flujo de datos de las operaciones, no los datos del tensor en sí. Por lo tanto, todos los cálculos deben realizarse utilizando otras operaciones de tensorflow, que se ejecutarán en el momento de la ejecución del gráfico.

Agregue sugerencias de tipo al registrar el gradiente personalizado para un tipo de operación para que el código sea más legible, depurable, más fácil de mantener y más sólido mediante la validación de datos. Por ejemplo, al tomar una op como parámetro en una función, especifique que la función de gradiente tomará una tf.Operation como tipo de parámetro.

Funciones de forma en C++

La API de TensorFlow tiene una función llamada "inferencia de formas" que proporciona información sobre las formas de los tensores sin tener que ejecutar el gráfico. La inferencia de formas está respaldada por "funciones de formas" que se registran para cada tipo de operación en la declaración REGISTER_OP de C++ y desempeñan dos funciones: afirmar que las formas de las entradas son compatibles durante la construcción del gráfico y especificar las formas para las salidas.

Las funciones de forma se definen como operaciones en la clase shape_inference::InferenceContext . Por ejemplo, en la función de forma para ZeroOut:

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

c->set_output(0, c->input(0)); declara que la forma de la primera salida debe establecerse en la forma de la primera entrada. Si la salida se selecciona por su índice como en el ejemplo anterior, el segundo parámetro de set_output debe ser un objeto ShapeHandle . Puede crear un objeto ShapeHandle vacío mediante su constructor predeterminado. El objeto ShapeHandle para una entrada con índice idx se puede obtener mediante c->input(idx) .

Hay una serie de funciones de forma comunes que se aplican a muchas operaciones, como shape_inference::UnchangedShape que se puede encontrar en common_shape_fns.h y se usa de la siguiente manera:

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

Una función de forma también puede restringir la forma de una entrada. Para la versión de ZeroOut con restricción de forma vectorial , la función de forma sería la siguiente:

    .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 llamada WithRank valida que la forma de entrada c->input(0) tiene una forma con exactamente una dimensión (o si la forma de entrada es desconocida, la forma de salida será un vector con una dimensión desconocida).

Si su operación es polimórfica con múltiples entradas , puede usar miembros de InferenceContext para determinar la cantidad de formas a verificar y Merge para validar que todas las formas sean compatibles (alternativamente, acceda a los atributos que indican las longitudes, con InferenceContext::GetAttr , que proporciona acceso a los atributos de la 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();
    });

Dado que la inferencia de formas es una característica opcional y las formas de los tensores pueden variar dinámicamente, las funciones de forma deben ser robustas ante la información de forma incompleta para cualquiera de las entradas. El método Merge en InferenceContext permite a la persona que llama afirmar que dos formas son iguales, incluso si una o ambas no tienen información completa. Las funciones de forma están definidas para todas las operaciones principales de TensorFlow y proporcionan muchos ejemplos de uso diferentes.

La clase InferenceContext tiene varias funciones que se pueden usar para definir manipulaciones de funciones de forma. Por ejemplo, puede validar que una dimensión particular tenga un valor muy específico usando InferenceContext::Dim e InferenceContext::WithValue ; puede especificar que una dimensión de salida sea la suma/producto de dos dimensiones de entrada usando InferenceContext::Add e InferenceContext::Multiply . Consulte la clase InferenceContext para conocer todas las manipulaciones de formas que puede especificar. El siguiente ejemplo establece la forma de la primera salida en (n, 3), donde la primera entrada tiene 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();
});

Si tiene una función de forma complicada, debería considerar agregar una prueba para validar que varias combinaciones de formas de entrada produzcan las combinaciones de formas de salida esperadas. Puede ver ejemplos de cómo escribir estas pruebas en algunas de nuestras pruebas de operaciones principales . (La sintaxis de INFER_OK e INFER_ERROR es un poco críptica, pero trate de ser compacta al representar las especificaciones de forma de entrada y salida en las pruebas. Por ahora, consulte los comentarios circundantes en esas pruebas para tener una idea de la especificación de la cadena de forma).

Cree un paquete pip para su operación personalizada

Para crear un paquete pip para su operación, consulte el ejemplo de tensorflow/custom-op . Esta guía muestra cómo crear operaciones personalizadas a partir del paquete pip de TensorFlow en lugar de crear TensorFlow desde el código fuente.