Ajuda a proteger a Grande Barreira de Corais com TensorFlow em Kaggle Junte Desafio

Criar uma operação

Se você deseja criar uma operação que não é coberta pela biblioteca TensorFlow existente, recomendamos que você primeiro tente escrever a operação em Python como uma composição de ops ou funções existentes do Python. Se isso não for possível, você pode criar um op C ++ personalizado. Existem vários motivos pelos quais você pode querer criar uma operação C ++ personalizada:

  • Não é fácil ou possível expressar sua operação como uma composição de operações existentes.
  • Não é eficiente expressar sua operação como uma composição de primitivas existentes.
  • Você deseja fundir manualmente uma composição de primitivas que um futuro compilador teria dificuldade em fundir.

Por exemplo, imagine que você deseja implementar algo como "pooling mediano", semelhante ao operador "MaxPool", mas calculando medianas em janelas deslizantes em vez de valores máximos. Pode ser possível fazer isso usando uma composição de operações (por exemplo, usando ExtractImagePatches e TopK), mas pode não ser tão eficiente em desempenho ou memória quanto uma operação nativa onde você pode fazer algo mais inteligente em uma única operação combinada. Como sempre, normalmente vale a pena primeiro tentar expressar o que você deseja usando a composição do operador, apenas escolhendo adicionar uma nova operação se isso se provar difícil ou ineficiente.

Para incorporar sua operação personalizada, você precisará:

  1. Registre a nova operação em um arquivo C ++. O registro de op define uma interface (especificação) para a funcionalidade do op, que é independente da implementação do op. Por exemplo, o registro do op define o nome do op e as entradas e saídas do op. Ele também define a função de forma que é usada para inferência de forma de tensor.
  2. Implemente o op em C ++. A implementação de um op é conhecida como kernel e é a implementação concreta da especificação registrada na Etapa 1. Pode haver vários kernels para diferentes tipos de entrada / saída ou arquiteturas (por exemplo, CPUs, GPUs).
  3. Crie um wrapper Python (opcional). Este wrapper é a API pública usada para criar o op em Python. Um wrapper padrão é gerado a partir do registro op, que pode ser usado diretamente ou adicionado.
  4. Escreva uma função para calcular gradientes para o op (opcional).
  5. Teste o op. Normalmente fazemos isso em Python por conveniência, mas você também pode testar o op em C ++. Se você definir gradientes, você pode verificá-los com o Python tf.test.compute_gradient_error . Ver relu_op_test.py como um exemplo que testa as funções para a frente de Relu semelhante operadores e seus gradientes.

Pré-requisitos

Defina a interface operacional

Você define a interface de uma operação registrando-a no sistema TensorFlow. No registro, você especificar o nome do seu op, suas entradas (tipos e nomes) e saídas (tipos e nomes), bem como docstrings e quaisquer attrs a op pode exigir.

Para ver como isso funciona, suponha que você gostaria de criar um op que leva um tensor de int32 s e envia uma cópia do tensor, com todos, mas o primeiro elemento set a zero. Para fazer isso, crie um arquivo chamado zero_out.cc . Em seguida, adicione uma chamada para o REGISTER_OP macro que define a interface para o seu 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();
    });

Este ZeroOut op tem um tensor to_zero de inteiros de 32 bits como entrada, e gera um tensor zeroed de inteiros de 32 bits. O op também usa uma função de forma para garantir que o tensor de saída tenha a mesma forma que o tensor de entrada. Por exemplo, se a entrada é um tensor de forma [10, 20], então esta função de forma especifica que a forma de saída também é [10, 20].

Implementar o kernel para o op

Depois de definir a interface, forneça uma ou mais implementações do op. Para criar um destes grãos, criar uma classe que se estende OpKernel e substitui o Compute método. O Compute método fornece um context argumento do tipo OpKernelContext* , a partir do qual você pode acessar coisas úteis como os tensores de entrada e saída.

Adicione seu kernel ao arquivo que você criou acima. O kernel pode ser parecido com isto:

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

Depois de implementar seu kernel, você o registra no sistema TensorFlow. No registro, você especifica diferentes restrições sob as quais este kernel será executado. Por exemplo, você pode ter um kernel feito para CPUs e outro separado para GPUs.

Para fazer isso para o ZeroOut op, adicione o seguinte ao zero_out.cc :

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

Kernels de CPU multi-threaded

Para escrever um kernel CPU multi-thread, a função Shard em work_sharder.h pode ser usado. Esta função fragmentos uma função de cálculo entre os fios configurados para ser utilizados para rosqueamento intra-operatório (ver intra_op_parallelism_threads em config.proto ).

Kernels GPU

Um kernel de GPU é implementado em duas partes: o OpKernel e o kernel CUDA e seu código de inicialização.

Às vezes, a implementação do OpKernel é comum entre um kernel de CPU e GPU, como na inspeção de entradas e na alocação de saídas. Nesse caso, uma implementação sugerida é:

  1. Defina o OpKernel modelado no dispositivo e o tipo primitivo do tensor.
  2. Para fazer o cálculo real da saída, a função Compute chama uma estrutura de functor modelo.
  3. A especialização desse functor para o CPUDevice é definida no mesmo arquivo, mas a especialização para o GPUDevice é definida em um arquivo .cu.cc, pois será compilado com o compilador CUDA.

Aqui está um exemplo de implementação.

// 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 a biblioteca de op

Compile a operação usando o compilador do sistema (instalação binária do TensorFlow)

Você deve ser capaz de compilar zero_out.cc com um C++ compilador tais como g++ ou clang disponível em seu sistema. O pacote binário PIP instala os arquivos de cabeçalho e a biblioteca de que você precisa para compilar seu op em locais que são específicos do sistema. No entanto, a biblioteca TensorFlow python fornece a get_include função para obter o diretório de cabeçalho, eo get_lib diretório tem um objeto compartilhado de ligação contra. Aqui estão as saídas dessas funções em uma 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'

Supondo que você tenha g++ instalado, aqui é a seqüência de comandos que você pode usar para compilar seu op em uma 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++11 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

No MacOS, o sinalizador adicional "dynamic_lookup -undefined" é necessária quando a construção do .so arquivo.

Nota sobre gcc versão >=5 : gcc usa o novo C ++ ABI desde a versão 5 . Os pacotes binários pip disponíveis no site da TensorFlow são construídos com gcc4 que usa a ABI mais velho. Se você compilar sua biblioteca op com gcc>=5 , adicione -D_GLIBCXX_USE_CXX11_ABI=0 para a linha de comando para tornar a biblioteca compatível com os abi mais velhos.

Compile a operação usando o bazel (instalação da fonte do TensorFlow)

Se você tem fontes do TensorFlow instaladas, pode usar o sistema de compilação do TensorFlow para compilar sua operação. Coloque um arquivo de construção com os seguintes Bazel regra de construção no tensorflow/core/user_ops diretório.

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

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

Execute o seguinte comando para construir zero_out.so .

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

Para compilar o Example operação, com o CUDA Kernel, você precisa usar os gpu_srcs parâmetro de tf_custom_op_library . Coloque um arquivo de construção com a seguinte regra de construção Bazel em uma nova pasta dentro do tensorflow/core/user_ops diretório (por exemplo, "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"],
)

Execute o seguinte comando para construir kernel_example.so .

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

Use o op em Python

TensorFlow Python API fornece a tf.load_op_library função para carregar a biblioteca dinâmica e registrar o op com o quadro TensorFlow. load_op_library retorna um módulo Python que contém os wrappers Python para a op eo kernel. Portanto, depois de construir o op, você pode fazer o seguinte para executá-lo no 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)

Tenha em mente, a função gerada será dado um nome snake_case (para cumprir com PEP8 ). Assim, se seu op é nomeado ZeroOut em arquivos do C ++, a função python será chamado zero_out .

Para fazer a op disponível como uma função regular import -able de um módulo Python, ele talvez útil ter o load_op_library chamada em um arquivo fonte Python da seguinte forma:

import tensorflow as tf

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

Verifique se a operação funciona

Uma boa maneira de verificar se você implementou com sucesso sua operação é escrever um teste para ela. Criar o arquivo zero_out_op_test.py com o conteúdo:

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

Em seguida, execute seu teste (supondo que você tenha o Tensorflow instalado):

$ python zero_out_op_test.py

Crie recursos avançados em sua operação

Agora que você sabe como construir uma operação e implementação básicas (e um tanto restritas), veremos algumas das coisas mais complicadas que você normalmente precisa para construir em sua operação. Isso inclui:

Verificações e validação condicionais

O exemplo acima assumiu que o op aplicado a um tensor de qualquer forma. E se fosse aplicado apenas a vetores? Isso significa adicionar uma verificação à implementação do OpKernel acima.

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

Isto assegura que a entrada é um vector, e retornos tendo definir o InvalidArgument estado se não for. O OP_REQUIRES macro recebe três argumentos:

Alternativamente, se você quiser testar se um Status objeto retornado de uma função é um erro, e se assim for devolvê-lo, use OP_REQUIRES_OK . Ambas as macros retornam da função em caso de erro.

Registro de op

Attrs

O Ops pode ter attrs, cujos valores são definidos quando o op é adicionado a um gráfico. Eles são usados ​​para configurar o op, e seus valores podem ser acessados ​​tanto na implementação do kernel quanto nos tipos de entradas e saídas no registro de op. Prefira usar uma entrada em vez de um atributo quando possível, uma vez que as entradas são mais flexíveis. Isso ocorre porque os atributos são constantes e devem ser definidos no momento da construção do gráfico. Em contraste, as entradas são tensores cujos valores podem ser dinâmicos; isto é, as entradas podem mudar a cada etapa, ser definidas usando um feed, etc. Attrs são usados ​​para coisas que não podem ser feitas com entradas: qualquer configuração que afeta a assinatura (número ou tipo de entradas ou saídas) ou que pode ' t mudar de etapa a etapa.

Você define um attr quando você registra a op, especificando seu nome e tipo usando o Attr método, que espera uma especificação da forma:

<name>: <attr-type-expr>

onde <name> começa com uma letra e pode ser composto por caracteres alfanuméricos e sublinhados, e <attr-type-expr> é um tipo de expressão da forma descrita abaixo .

Por exemplo, se você gostaria que o ZeroOut op para preservar um índice especificado pelo usuário, em vez de apenas o elemento 0, você pode registrar o op assim:

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

(Note-se que o conjunto de tipos de atributos é diferente do tf.DType utilizado para entradas e saídas).

Seu kernel pode então acessar esta attr em seu construtor através do context parâmetro:

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 pode então ser utilizado no Compute método:

  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 Atr

Os seguintes tipos são suportados em um atributo:

  • string : Qualquer sequência de bytes (não necessitam de estar UTF-8).
  • int : um inteiro assinado.
  • float : Um número de ponto flutuante.
  • bool : Verdadeiro ou falso.
  • type : Um dos (não-ref) valores de DataType .
  • shape : Uma TensorShapeProto .
  • list(<type>) : Uma lista de <type> , onde <type> é um dos tipos acima. Note-se que list(list(<type>)) é inválido.

Veja também: op_def_builder.cc:FinalizeAttr para uma lista definitiva.

Valores padrão e restrições

Attrs pode ter valores padrão e alguns tipos de attrs podem ter restrições. Para definir um attr com restrições, você pode usar o seguinte <attr-type-expr> s:

{'<string1>', '<string2>'} : O valor deve ser uma cadeia que tem ou o valor <string1> ou <string2> . O nome do tipo, string , está implícito quando você usa esta sintaxe. Isso emula um enum:

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

{<type1>, <type2>} : O valor é do tipo type , e deve ser um dos <type1> ou <type2> , onde <type1> e <type2> são suportados tf.DType . Você não especificar que o tipo do attr é type . Isso está implícito quando você tem uma lista de tipos de {...} . Por exemplo, neste caso o attr t é um tipo que tem de ser um int32 , um float , ou um bool :

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

Existem atalhos para restrições de tipo comuns:

  • numbertype : Tipo type restrito ao numérico (não-string e não-bool) tipos.
  • realnumbertype : Como numbertype sem tipos complexos.
  • quantizedtype : Como numbertype mas apenas os tipos de números quantizados.

As listas específicas de tipos permitidos por estes são definidos pelas funções (como NumberTypes() ) em tensorflow/core/framework/types.h . Neste exemplo, o attr t deve ser um dos tipos numéricos:

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

Para esta operação:

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

As listas podem ser combinadas com outras listas e tipos únicos. O seguinte op permite attr t de ser qualquer um dos tipos numéricos, ou do tipo booleano:

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

Para esta operação:

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> : O valor deve ser um inteiro cujo valor é maior do que ou igual a <n> , onde <n> é um número natural. Por exemplo, o seguinte especifica registro op que o attr a deve ter um valor que é pelo menos 2 :

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

list(<type>) >= <n> : Uma lista de tipo <type> cujo comprimento é maior do que ou igual a <n> . Por exemplo, o seguinte especifica de registo op que o attr a é uma lista de tipos (quer int32 ou float ), e que deve haver pelo menos três delas:

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

Para definir um valor padrão para um attr (tornando-se opcional no código gerado), adicione = <default> até ao fim, como em:

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

Além disso, uma restrição e um valor padrão podem ser especificados:

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

A sintaxe com suporte do valor padrão é o que seria usado na representação proto da definição GraphDef resultante.

Aqui estão alguns exemplos de como especificar um padrão para todos os 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]");

Nota, em particular, que os valores do tipo type uso tf.DType .

Polimorfismo

Polimorfismo de tipo

Para ops que pode tomar diferentes tipos como entrada ou produzem diferentes tipos de saída, você pode especificar um attr em um tipo de entrada ou de saída no registro op. Normalmente você iria em seguida, registrar uma OpKernel para cada tipo suportado.

Por exemplo, se você gostaria que o ZeroOut op ao trabalho na float s além de int32 s, seu registro op pode parecer:

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

O seu registo op agora especifica que o tipo da entrada deve ser float , ou int32 , e que sua saída vai ser do mesmo tipo, já que ambos têm tipo T .

Nomeação

Entradas, saídas e atributos geralmente devem receber nomes snake_case. A única exceção são os atributos que são usados ​​como o tipo de uma entrada ou o tipo de uma saída. Esses atributos podem ser inferidos quando o op é adicionado ao gráfico e, portanto, não aparecem na função do op. Por exemplo, esta última definição de ZeroOut irá gerar uma função Python que se parece com:

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 é passado um int32 tensor, então T é automaticamente definida para int32 (bem, na verdade DT_INT32 ). Esses atributos inferidos recebem nomes em maiúsculas ou CamelCase.

Compare isso com um op que tem um tipo attr que determina o tipo de saída:

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

Nesse caso, o usuário deve especificar o tipo de saída, como no Python gerado:

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`.
  """
Exemplo 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 a compatibilidade com versões anteriores , você deve especificar um valor padrão ao adicionar um attr para um op existente:

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

Vamos dizer que você queria adicionar mais tipos, dizem double :

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

Em vez de escrever outro OpKernel com código redundante como acima, muitas vezes você vai ser capaz de usar um modelo de C ++ em seu lugar. Você ainda terá um registro de kernel ( REGISTER_KERNEL_BUILDER chamada) 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>);

Se você tiver mais de algumas sobrecargas, poderá colocar o registro em uma 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

Dependendo da lista de tipos que você está registrando o kernel para, você pode ser capaz de usar uma macro fornecida pelo 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 e saídas

Além de serem capazes de aceitar ou produzir diferentes tipos, os ops podem consumir ou produzir um número variável de tensores.

No seguinte exemplo, o attr T mantém uma lista de tipos, e é usado como o tipo de tanto a entrada in e a saída out . A entrada e saída são listas de tensores de que tipo (e o número e tipos de tensores na saída têm o mesmo que o de entrada, uma vez que ambos têm tipo T ).

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

Você também pode colocar restrições sobre os tipos que podem ser especificados na lista. Neste caso seguinte, a entrada é uma lista de float e double tensores. O op aceita, por exemplo, os tipos de entrada (float, double, float) e, nesse caso, o tipo de saída também seria (float, double, float) .

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

Se você quiser que todos os tensores em uma lista sejam do mesmo tipo, você pode fazer algo como:

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

Este aceita uma lista de int32 tensores, e usa um int attr N para especificar o tamanho da lista.

Isto pode ser feito tipo polimórfico bem. No seguinte exemplo, a entrada é uma lista dos tensores (com um comprimento de "N" ) do mesmo (mas não especificado) tipo ( "T" ), e a saída é um único tipo de tensor de correspondência:

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

Por padrão, o tensor listas têm um comprimento mínimo de 1. Você pode mudar esse padrão usando um ">=" restrição na attr correspondente . No exemplo a seguir, a entrada é uma lista de pelo menos 2 int32 tensores:

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

A mesma sintaxe funciona com "list(type)" attrs:

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

Entradas e saídas

Para resumir o acima, um registro de operação pode ter várias entradas e saídas:

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

Cada especificação de entrada ou saída tem o formato:

<name>: <io-type-expr>

onde <name> começa com uma letra e pode ser composto de caracteres alfanuméricos e sublinhados. <io-type-expr> é uma das seguintes expressões do tipo:

  • <type> , onde <type> é um tipo de entrada suportado (por exemplo, float , int32 , string ). Isso especifica um único tensor do tipo fornecido.

    Veja tf.DType .

    REGISTER_OP("BuiltInTypesExample")
        .Input("integers: int32")
        .Input("complex_numbers: complex64");
    
  • <attr-type> , onde <attr-type> é o nome de um Atr com o tipo de type ou list(type) (com um possível restrição tipo). Esta sintaxe permite ops polimórficas .

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

    Fazendo referência a um attr do tipo list(type) permite-lhe aceitar uma seqüência 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");
    

    Note-se que o número e tipos de tensores na saída out é o mesmo que na entrada in , uma vez que ambos são do tipo T .

  • Para uma sequência de tensores com o mesmo tipo: <number> * <type> , onde <number> é o nome de um Atr com tipo int . O <type> pode ser um tf.DType , ou o nome de um attr com o tipo de type . Como exemplo do primeiro, este op aceita uma lista de int32 tensores:

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

    Considerando que esta op aceita uma lista de tensores de qualquer tipo, desde que sejam todos iguais:

    REGISTER_OP("SameTypeSequenceExample")
        .Attr("NumTensors: int")
        .Attr("T: type")
        .Input("in: NumTensors * T")
    
  • Para uma referência a um tensor: Ref(<type>) , onde <type> é um dos tipos anteriores.

Qualquer atributo usado no tipo de entrada será inferido. Por convenção essas attrs inferidas usar nomes de capital (como T ou N ). Caso contrário, entradas, saídas e attrs têm nomes como parâmetros de função (por exemplo, num_outputs ). Para mais detalhes, consulte a seção anterior sobre nomeação .

Para mais detalhes, consulte tensorflow/core/framework/op_def_builder.h .

Compatibilidade com versões anteriores

Vamos supor que você escreveu uma boa operação personalizada e compartilhou com outras pessoas, portanto, você tem clientes satisfeitos usando sua operação. No entanto, você gostaria de fazer alterações na operação de alguma forma.

Em geral, muda para existente, o check-in especificações deve ser compatível com versões anteriores: mudando a especificação de um op não deve quebrar serializados anteriores GraphDef buffers de protocolo construídos a partir de especificações mais velhos. Os detalhes do GraphDef compatibilidade são descritos aqui .

Existem várias maneiras de preservar a compatibilidade com versões anteriores.

  1. Qualquer novo atributo adicionado a uma operação deve ter valores padrão definidos e, com esse valor padrão, o op deve ter o comportamento original. Para alterar uma operação de não polimórficos a polimórfica, você deve dar um valor padrão para o novo tipo attr para preservar a assinatura original por padrão. Por exemplo, se sua operação foi:

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

    você pode torná-lo polimórfico de maneira compatível com versões anteriores usando:

    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: T")
        .Output("out: T")
        .Attr("T: numerictype = DT_FLOAT");
    
  2. Você pode tornar uma restrição em um atrativo menos restritiva com segurança. Por exemplo, você pode mudar de {int32, int64} para {int32, int64, float} ou type . Ou você pode mudar de {"apple", "orange"} para {"apple", "banana", "orange"} ou string .

  3. Você pode alterar entradas / saídas únicas em entradas / saídas de lista, desde que o padrão para o tipo de lista corresponda à assinatura antiga.

  4. Você pode adicionar uma nova lista de entrada / saída, se o padrão for vazio.

  5. Crie um namespace para quaisquer novos ops que você criar, prefixando os nomes dos ops com algo exclusivo para o seu projeto. Isso evita que sua operação colida com qualquer operação que possa ser incluída em versões futuras do TensorFlow.

  6. Planejar com antecedência! Tente antecipar usos futuros para a operação. Algumas alterações de assinatura não podem ser feitas de maneira compatível (por exemplo, transformar uma lista do mesmo tipo em uma lista de vários tipos).

A lista completa de mudanças seguras e inseguras podem ser encontrados em tensorflow/core/framework/op_compatibility_test.cc . Se você não puder fazer sua mudança para uma operação compatível com versões anteriores, crie uma nova operação com um novo nome com a nova semântica.

Observe também que, enquanto essas mudanças podem manter GraphDef compatibilidade, o código Python gerado pode mudar de uma forma que não é compatível com o interlocutor idade. A API Python pode ser mantida compatível por meio de mudanças cuidadosas em um wrapper Python escrito à mão, mantendo a assinatura antiga, exceto possivelmente adicionando novos argumentos opcionais ao final. Geralmente mudanças incompatíveis só pode ser feita quando TensorFlow muda as versões principais, e deve estar de acordo com as GraphDef versão semântica .

Suporte para GPU

Você pode implementar diferentes OpKernels e registrar um para CPU e outro para GPU, assim como você pode registrar kernels para diferentes tipos . Existem vários exemplos de grãos com suporte GPU em tensorflow/core/kernels/ . Observe alguns grãos têm uma versão CPU em um .cc arquivo, uma versão GPU em um arquivo terminando em _gpu.cu.cc , e alguns código compartilhado em comum em um .h arquivo.

Por exemplo, o tf.pad tem tudo, mas a GPU kernel no tensorflow/core/kernels/pad_op.cc . A GPU kernel está em tensorflow/core/kernels/pad_op_gpu.cu.cc , eo código compartilhado é uma classe de modelo definido na tensorflow/core/kernels/pad_op.h . Organizamos o código dessa maneira por dois motivos: permite que você compartilhe um código comum entre as implementações de CPU e GPU e coloca a implementação da GPU em um arquivo separado para que possa ser compilado apenas pelo compilador da GPU.

Uma coisa a nota, mesmo quando a versão do kernel GPU da pad é usado, ele ainda precisa de seu "paddings" de entrada na memória CPU. Para marcar esse entradas ou saídas são mantidos na CPU, adicione um HostMemory() chamada para o registro de kernel, por exemplo:

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

Compilar o kernel para o dispositivo GPU

Olhada cuda_op_kernel.cu.cc para um exemplo que usa um kernel CUDA para implementar um op. O tf_custom_op_library aceita um gpu_srcs argumento em que a lista de arquivos de origem contendo os kernels CUDA ( *.cu.cc arquivos) pode ser especificado. Para uso com uma instalação binária de TensorFlow, os kernels CUDA tem que ser compilado com NVIDIA nvcc compilador. Aqui está a seqüência de comandos que você pode usar para compilar o cuda_op_kernel.cu.cc e cuda_op_kernel.cc em uma única biblioteca dinâmica:

nvcc -std=c++11 -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++11 -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 produzido acima, pode ser carregado como de costume em Python, usando o tf.load_op_library função.

Nota que, se as suas bibliotecas CUDA não estão instalados em /usr/local/lib64 , você precisa especificar o caminho explicitamente no comando (g ++) acima segundo. Por exemplo, adicione -L /usr/local/cuda-8.0/lib64/ se o seu CUDA está instalado no /usr/local/cuda-8.0 .

Implementar o gradiente em Python

Dado um gráfico de operações, o TensorFlow usa diferenciação automática (retropropagação) para adicionar novas operações que representam gradientes em relação às operações existentes. Para fazer a diferenciação automática funcionar para novos ops, você deve registrar uma função gradiente que calcula gradientes em relação às entradas dos ops, dados gradientes em relação às saídas dos ops.

Matematicamente, se calcula um op \(y = f(x)\) os registados convertidos gradiente op gradientes \(\partial L/ \partial y\) de perda \(L\) com respeito a\(y\) em gradientes \(\partial L/ \partial x\) com respeito a \(x\) através da regra da cadeia:

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

No caso de ZeroOut , apenas uma entrada na entrada afecta a saída, de modo que o gradiente com respeito à entrada é um esparso "um quente" tensor. Isso é expresso da seguinte forma:

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

Os detalhes sobre registrar funções de gradiente com tf.RegisterGradient :

  • Para um op com uma saída, a função de gradiente terá um tf.Operation , op , e um tf.Tensor grad e construir novas operações para fora dos tensores op.inputs[i] , op.outputs[i] , e grad . Informações sobre quaisquer attrs podem ser encontrados através tf.Operation.get_attr .

  • Se o op tem múltiplas saídas, a função de gradiente terá op e grads , onde grads é uma lista de gradientes com respeito a cada saída. O resultado da função de gradiente deve ser uma lista de Tensor objectos que representam os gradientes com respeito a cada entrada.

  • Se não é gradiente não bem definido para alguma entrada, tal como para entradas inteiros usados como índices, o gradiente devolvido correspondente deve ser None . Por exemplo, para um op tendo um ponto flutuante tensor x e um índice inteiro i , a função de gradiente iria return [x_grad, None] .

  • Se não houver gradiente significativo para a operação, geralmente você não terá que registrar qualquer gradiente e, desde que o gradiente da operação nunca seja necessário, você ficará bem. Em alguns casos, um op não tem gradiente bem definido, mas pode estar envolvido no cálculo do gradiente. Aqui você pode usar ops.NotDifferentiable a zeros para trás automaticamente propagar.

Observe que, no momento em que a função gradiente é chamada, apenas o gráfico de fluxo de dados de ops está disponível, não os dados do tensor em si. Assim, todos os cálculos devem ser realizados usando outras operações de fluxo de tensor, para serem executados no tempo de execução do gráfico.

Funções de forma em C ++

A API TensorFlow tem um recurso chamado "inferência de formas", que fornece informações sobre as formas dos tensores sem a necessidade de executar o gráfico. Forma inferência é suportada por "funções de forma" que são registadas para cada tipo op no C ++ REGISTER_OP declaração, e realizam duas funções: afirmar que as formas das entradas são compatíveis durante a construção do gráfico, e que especificam as formas para as saídas.

Funções de forma são definidos como operações no shape_inference::InferenceContext classe. Por exemplo, na função 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 a forma da primeira saída deve ser definida para a forma da primeira entrada. Se a saída é seleccionado pelo seu índice como no exemplo acima, o segundo parâmetro de set_output deve ser um ShapeHandle objecto. Você pode criar um vazio ShapeHandle objeto pelo seu construtor padrão. O ShapeHandle objecto para uma entrada com índice idx pode ser obtido por c->input(idx) .

Há um número de funções de forma comuns que se aplicam a várias operações, tais como shape_inference::UnchangedShape que pode ser encontrado em common_shape_fns.h e utilizados como se segue:

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

Uma função de forma também pode restringir a forma de uma entrada. Para a versão de ZeroOut com um constrangimento de forma vector , a função de forma seria como se segue:

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

Os WithRank valida chamada que a forma de entrada c->input(0) tem uma forma com exatamente uma dimensão (ou se a forma de entrada é desconhecida, a forma de saída será um vetor com uma dimensão desconhecida).

Se o seu op é polimórfico com múltiplas entradas , você pode usar os membros da InferenceContext para determinar o número de formas de verificar e Merge para validar que as formas são compatíveis (alternativamente, atributos de acesso que indicam os comprimentos, com InferenceContext::GetAttr , que fornece acesso aos atributos da 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();
    });

Uma vez que a inferência de forma é um recurso opcional e as formas dos tensores podem variar dinamicamente, as funções de forma devem ser robustas para informações de forma incompletas para qualquer uma das entradas. A Merge método em InferenceContext permite que o chamador para afirmar que duas formas são as mesmas, mesmo que um ou ambos deles ainda não temos informação completa. As funções de forma são definidas para todas as operações principais do TensorFlow e fornecem muitos exemplos de uso diferentes.

O InferenceContext classe tem um número de funções que podem ser utilizadas para definir a forma da função manipulações. Por exemplo, você pode validar que uma dimensão particular tem um valor muito específico usando InferenceContext::Dim e InferenceContext::WithValue ; você pode especificar que uma dimensão de saída é a soma / produto de duas dimensões de entrada usando InferenceContext::Add e InferenceContext::Multiply . Veja a InferenceContext classe para todas as várias manipulações de forma que você pode especificar. O exemplo a seguir define a forma da primeira saída para (n, 3), onde a primeira entrada tem a 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 você tiver uma função de forma complicada, deve considerar adicionar um teste para validar se várias combinações de forma de entrada produzem as combinações de formas de saída esperadas. Você pode ver exemplos de como escrever esses testes em alguns nossos testes ops núcleo . (A sintaxe de INFER_OK e INFER_ERROR são um pouco enigmática, mas tente ser compacta na representação especificações de entrada e saída de forma em testes. Por enquanto, ver os comentários que cercam nesses testes para ter uma noção da especificação corda forma).

Crie um pacote pip para sua operação personalizada

Para construir um pip pacote para o seu op, consulte o tensorflow / custom-op exemplo. Este guia mostra como criar operações personalizadas a partir do pacote TensorFlow pip em vez de criar TensorFlow a partir da fonte.