Благодарим за настройку Google I/O. Посмотреть все сеансы по запросу Смотреть по запросу

Создать операцию

Если вы хотите создать операцию, которая не поддерживается существующей библиотекой TensorFlow, мы рекомендуем вам сначала попробовать написать операцию на Python как композицию существующих операций или функций Python. Если это невозможно, вы можете создать пользовательскую операцию C++. Есть несколько причин, по которым вы можете захотеть создать пользовательскую операцию C++:

  • Нелегко или невозможно выразить вашу операцию как композицию существующих операций.
  • Неэффективно выражать вашу операцию как композицию существующих примитивов.
  • Вы хотите вручную слить композицию примитивов, слияние которой будущему компилятору будет затруднительно.

Например, представьте, что вы хотите реализовать что-то вроде "медианного пула", аналогичного оператору "MaxPool", но вычисляющего медианы по скользящим окнам вместо максимальных значений. Выполнение этого с использованием композиции операций может быть возможным (например, с помощью ExtractImagePatches и TopK), но может быть не таким эффективным с точки зрения производительности или памяти, как собственная операция, когда вы можете сделать что-то более умное в одной объединенной операции. Как всегда, обычно сначала стоит попытаться выразить то, что вы хотите, с помощью композиции операторов, а добавлять новую операцию только в том случае, если это окажется сложным или неэффективным.

Чтобы включить свою пользовательскую операцию, вам необходимо:

  1. Зарегистрируйте новую операцию в файле C++. Регистрация оператора определяет интерфейс (спецификацию) для функциональности оператора, который не зависит от реализации оператора. Например, регистрация операции определяет имя операции, а также ее входы и выходы. Он также определяет функцию формы, которая используется для вывода формы тензора.
  2. Реализуйте операцию на C++. Реализация операции известна как ядро, и это конкретная реализация спецификации, которую вы зарегистрировали на шаге 1. Может быть несколько ядер для разных типов ввода/вывода или архитектур (например, ЦП, ГП).
  3. Создайте оболочку Python (необязательно). Эта оболочка представляет собой общедоступный API, который используется для создания операции в Python. Оболочка по умолчанию создается из регистрации операции, которую можно использовать напрямую или добавить к ней.
  4. Напишите функцию для вычисления градиентов для операции (необязательно).
  5. Протестируйте оп. Обычно для удобства мы делаем это на Python, но вы также можете протестировать операцию на C++. Если вы определяете градиенты, вы можете проверить их с помощью Python tf.test.compute_gradient_error . См. relu_op_test.py в качестве примера, который проверяет прямые функции Relu-подобных операторов и их градиенты.

Предпосылки

Определить операционный интерфейс

Вы определяете интерфейс оператора, регистрируя его в системе TensorFlow. При регистрации вы указываете имя своей операции, ее входные данные (типы и имена) и выходные данные (типы и имена), а также строки документации и любые атрибуты, которые могут потребоваться операции.

Чтобы увидеть, как это работает, предположим, что вы хотите создать операцию, которая принимает тензор из int32 и выводит копию тензора со всеми элементами, кроме первого, равными нулю. Для этого создайте файл с именем zero_out.cc . Затем добавьте вызов макроса REGISTER_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();
    });

Эта операция ZeroOut принимает на вход один тензор to_zero из 32-битных целых чисел и выводит zeroed тензор из 32-битных целых чисел. Оператор также использует функцию формы, чтобы гарантировать, что выходной тензор имеет ту же форму, что и входной тензор. Например, если вход представляет собой тензор формы [10, 20], то эта функция формы указывает, что выходная форма также [10, 20].

Реализовать ядро ​​для операции

После определения интерфейса предоставьте одну или несколько реализаций оператора op. Чтобы создать одно из этих ядер, создайте класс, расширяющий OpKernel и переопределяющий метод Compute . Метод Compute предоставляет один аргумент context типа OpKernelContext* , из которого вы можете получить доступ к таким полезным вещам, как входные и выходные тензоры.

Добавьте ваше ядро ​​в файл, который вы создали выше. Ядро может выглядеть примерно так:

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

После реализации вашего ядра вы регистрируете его в системе TensorFlow. При регистрации вы указываете различные ограничения, при которых будет работать это ядро. Например, у вас может быть одно ядро ​​для ЦП и отдельное для ГП.

Чтобы сделать это для операции ZeroOut , добавьте в zero_out.cc следующее:

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

Многопоточные ядра ЦП

Для написания многопоточного ядра ЦП можно использовать функцию Shard в work_sharder.h . Эта функция разбивает вычислительную функцию на потоки, сконфигурированные для использования для потоков внутри операции (см. intra_op_parallelism_threads в config.proto ).

Ядра графического процессора

Ядро GPU состоит из двух частей: ядра OpKernel и ядра CUDA и кода его запуска.

Иногда реализация OpKernel является общей для ядра CPU и GPU, например, при проверке входных данных и распределении выходных данных. В этом случае предлагаемая реализация заключается в следующем:

  1. Определите шаблон OpKernel на устройстве и примитивный тип тензора.
  2. Чтобы выполнить фактическое вычисление выходных данных, функция Compute вызывает шаблонную структуру функтора.
  3. Специализация этого функтора для CPUDevice определяется в том же файле, но специализация для GPUDevice определяется в файле .cu.cc, поскольку он будет скомпилирован с помощью компилятора CUDA.

Вот пример реализации.

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

Создайте операционную библиотеку

Скомпилируйте операцию с помощью системного компилятора (бинарная установка TensorFlow)

Вы должны быть в состоянии скомпилировать zero_out.cc с помощью компилятора C++ , такого как g++ или clang доступного в вашей системе. Двоичный пакет PIP устанавливает файлы заголовков и библиотеку, необходимые для компиляции вашей операции, в местах, зависящих от системы. Однако библиотека Python TensorFlow предоставляет функцию get_include для получения каталога заголовков, а каталог get_lib имеет общий объект для ссылки. Вот результаты этих функций на компьютере с 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'

Предполагая, что у вас установлен g++ , вот последовательность команд, которые вы можете использовать для компиляции вашей операции в динамическую библиотеку.

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

В macOS при создании файла .so требуется дополнительный флаг «-undefined dynamic_lookup».

Обратите внимание на версию gcc >=5 : начиная с версии 5 , gcc использует новый C++ ABI . TensorFlow 2.8 и более ранние версии были созданы с помощью gcc4 , использующего старый ABI. Если вы используете эти версии TensorFlow и пытаетесь скомпилировать свою операционную библиотеку с gcc>=5 , добавьте -D_GLIBCXX_USE_CXX11_ABI=0 в командную строку, чтобы сделать библиотеку совместимой со старым ABI. Пакеты TensorFlow 2.9+ по умолчанию совместимы с новым ABI.

Скомпилируйте операцию с помощью bazel (установка исходного кода TensorFlow)

Если у вас установлены исходники TensorFlow, вы можете использовать систему сборки TensorFlow для компиляции своей операции. Поместите файл BUILD со следующим правилом сборки Bazel в каталог tensorflow/core/user_ops .

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

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

Выполните следующую команду, чтобы построить zero_out.so .

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

Для компиляции операции Example с ядром CUDA необходимо использовать параметр gpu_srcs для tf_custom_op_library . Поместите файл BUILD со следующим правилом сборки Bazel в новую папку внутри каталога tensorflow/core/user_ops (например, «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"],
)

Выполните следующую команду, чтобы собрать kernel_example.so .

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

Используйте операцию в Python

TensorFlow Python API предоставляет функцию tf.load_op_library для загрузки динамической библиотеки и регистрации операции в среде TensorFlow. load_op_library возвращает модуль Python, содержащий оболочки Python для операции и ядра. Таким образом, после того, как вы создали операцию, вы можете сделать следующее, чтобы запустить ее из 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)

Имейте в виду, что сгенерированной функции будет присвоено имя snake_case (для соответствия PEP8 ). Итак, если ваша операция называется ZeroOut в файлах C++, функция python будет называться zero_out .

Чтобы сделать операцию доступной как обычную функцию, import из модуля Python, может быть полезно иметь вызов load_op_library в исходном файле Python следующим образом:

import tensorflow as tf

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

Убедитесь, что оп работает

Хороший способ убедиться, что вы успешно реализовали свою операцию, — написать для нее тест. Создайте файл zero_out_op_test.py с содержимым:

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

Затем запустите тест (при условии, что у вас установлен тензорный поток):

$ python zero_out_op_test.py

Внедряйте расширенные функции в свою операцию

Теперь, когда вы знаете, как построить базовую (и несколько ограниченную) операцию и реализацию, мы рассмотрим некоторые из более сложных вещей, которые вам обычно нужно будет встроить в вашу операцию. Это включает в себя:

Условные проверки и валидация

В приведенном выше примере предполагалось, что операция применяется к тензору любой формы. Что, если бы это относилось только к векторам? Это означает добавление проверки к приведенной выше реализации OpKernel.

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

Это подтверждает, что ввод является вектором, и возвращает значение InvalidArgument , если это не так. Макрос OP_REQUIRES принимает три аргумента:

  • context , который может быть указателем OpKernelContext или OpKernelConstruction (см. tensorflow/core/framework/op_kernel.h ), для его метода SetStatus() .
  • Состояние. Например, в tensorflow/core/framework/tensor_shape.h есть функции для проверки формы тензора.
  • Сама ошибка, представленная объектом Status , см. в tensorflow/core/platform/status.h . Status имеет как тип (часто InvalidArgument , но см. список типов), так и сообщение. Функции для построения ошибки можно найти в tensorflow/core/platform/errors.h .

В качестве альтернативы, если вы хотите проверить, является ли объект Status , возвращенный какой-либо функцией, ошибкой, и если да, то верните его, используйте OP_REQUIRES_OK . Оба этих макроса возвращаются из функции при ошибке.

Оп регистрация

атрибуты

У операций могут быть атрибуты, значения которых устанавливаются при добавлении операции в граф. Они используются для настройки операции, и к их значениям можно получить доступ как в реализации ядра, так и в типах входных и выходных данных при регистрации операции. По возможности предпочтительнее использовать ввод вместо атрибута, поскольку ввод более гибок. Это связано с тем, что атрибуты являются константами и должны быть определены во время построения графа. Напротив, входные данные представляют собой тензоры, значения которых могут быть динамическими; то есть входы могут меняться на каждом шагу, задаваться с помощью фида и т. д. Attrs используются для вещей, которые нельзя сделать с входами: любая конфигурация, которая влияет на сигнатуру (количество или тип входов или выходов) или которая может' t меняется от шага к шагу.

Вы определяете атрибут при регистрации оператора, указывая его имя и тип с помощью метода Attr , который ожидает спецификацию формы:

<name>: <attr-type-expr>

где <name> начинается с буквы и может состоять из буквенно-цифровых символов и знаков подчеркивания, а <attr-type-expr> — это выражение типа в форме, описанной ниже .

Например, если вы хотите, чтобы операция ZeroOut сохраняла указанный пользователем индекс, а не только 0-й элемент, вы можете зарегистрировать операцию следующим образом:

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

(Обратите внимание, что набор типов атрибутов отличается от tf.DType , используемого для входных и выходных данных.)

Затем ваше ядро ​​​​может получить доступ к этому атрибуту в своем конструкторе через параметр 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_;
};

который затем можно использовать в методе 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_);
  }

Типы атрибутов

Следующие типы поддерживаются в attr:

  • string : любая последовательность байтов (не обязательно UTF8).
  • int : Целое число со знаком.
  • float : число с плавающей запятой.
  • bool : Истинно или ложно.
  • type : одно из (не ref) значений DataType .
  • shape : TensorShapeProto .
  • list(<type>) : список <type> , где <type> — один из указанных выше типов. Обратите внимание, что list(list(<type>)) недействителен.

См. также: op_def_builder.cc:FinalizeAttr для окончательного списка.

Значения по умолчанию и ограничения

Attrs могут иметь значения по умолчанию, а некоторые типы attrs могут иметь ограничения. Чтобы определить attr с ограничениями, вы можете использовать следующие <attr-type-expr> s:

{'<string1>', '<string2>'} : значение должно быть строкой со значением <string1> или <string2> . Имя типа string подразумевается при использовании этого синтаксиса. Это эмулирует перечисление:

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

{<type1>, <type2>} : значение имеет тип type и должно быть одним из <type1> или <type2> , где <type1> и <type2> поддерживаются tf.DType . Вы не указываете тип атрибута type . Это подразумевается, когда у вас есть список типов в {...} . Например, в этом случае attr t — это тип, который должен быть int32 , float или bool :

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

Существуют ярлыки для ограничений общего типа:

  • numbertype : Тип type ограничен числовыми (не строковыми и не логическими) типами.
  • realnumbertype : как numbertype без сложных типов.
  • quantizedtype : как numbertype , но только типы квантованных чисел.

Конкретные списки типов, разрешенных ими, определяются функциями (например NumberTypes() ) в tensorflow/core/framework/types.h . В этом примере attr t должен быть одним из числовых типов:

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

Для этой операции:

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

Списки можно комбинировать с другими списками и отдельными типами. Следующая операция позволяет attr t быть любого числового типа или типа bool:

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

Для этой операции:

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> : значение должно быть целым числом, значение которого больше или равно <n> , где <n> — натуральное число. Например, следующая регистрация операции указывает, что атрибут a должен иметь значение не менее 2 :

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

list(<type>) >= <n> : список типа <type> , длина которого больше или равна <n> . Например, следующая регистрация op указывает, что attr a представляет собой список типов (либо int32 , либо float ), и что их должно быть не менее 3:

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

Чтобы установить значение по умолчанию для атрибута (сделав его необязательным в сгенерированном коде), добавьте = <default> в конец, например:

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

Кроме того, можно указать как ограничение, так и значение по умолчанию:

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

Поддерживаемый синтаксис значения по умолчанию — это то, что будет использоваться в прото-представлении результирующего определения GraphDef.

Вот примеры того, как указать значение по умолчанию для всех типов:

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

Обратите внимание, в частности, что значения типа type используют tf.DType .

Полиморфизм

Полиморфизм типов

Для операций, которые могут принимать разные типы в качестве входных данных или создавать разные типы выходных данных, вы можете указать атрибут во входном или выходном типе при регистрации операции. Обычно вы должны зарегистрировать OpKernel для каждого поддерживаемого типа.

Например, если вы хотите, чтобы операция ZeroOut работала с float в дополнение к int32 , ваша регистрация операции может выглядеть так:

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

Ваша регистрация операции теперь указывает, что тип ввода должен быть float или int32 , и что его вывод будет того же типа, поскольку оба имеют тип T

Именование

Входы, выходы и атрибуты обычно должны иметь имена в змеином регистре. Единственным исключением являются атрибуты, которые используются в качестве типа ввода или в качестве типа вывода. Эти атрибуты можно вывести, когда оператор добавляется к графу, и поэтому они не отображаются в функции оператора. Например, это последнее определение ZeroOut сгенерирует функцию Python, которая выглядит так:

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

Если to_zero передается тензор int32 , то T автоматически устанавливается в int32 (ну, на самом деле DT_INT32 ). Эти предполагаемые атрибуты получают имена с заглавными буквами или CamelCase.

Сравните это с операцией, которая имеет атрибут типа, определяющий тип вывода:

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

В этом случае пользователь должен указать тип вывода, как в сгенерированном Python:

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`.
  """
Пример полиморфизма типов
#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);

Чтобы сохранить обратную совместимость , вы должны указать значение по умолчанию при добавлении атрибута к существующей операции:

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

Допустим, вы хотите добавить больше типов, скажем, double :

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

Вместо написания еще одного OpKernel с избыточным кодом, как указано выше, часто можно использовать шаблон C++. У вас по-прежнему будет одна регистрация ядра (вызов REGISTER_KERNEL_BUILDER ) для каждой перегрузки.

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

Если у вас больше пары перегрузок, вы можете поместить регистрацию в макрос.

#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

В зависимости от списка типов, для которых вы регистрируете ядро, вы можете использовать макрос, предоставленный 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
Список входов и выходов

В дополнение к способности принимать или создавать различные типы, операции могут потреблять или создавать переменное количество тензоров.

В следующем примере attr T содержит список типов и используется в качестве типа как для in так и для out . Вход и выход — это списки тензоров этого типа (и количество и типы тензоров на выходе такие же, как и на входе, поскольку оба имеют тип T ).

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

Вы также можете установить ограничения на то, какие типы могут быть указаны в списке. В следующем случае входными данными является список тензоров с float и double чисел. Оператор принимает, например, типы ввода (float, double, float) , и в этом случае тип вывода также будет (float, double, float) .

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

Если вы хотите, чтобы все тензоры в списке были одного типа, вы можете сделать что-то вроде:

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

Это принимает список тензоров int32 и использует int attr N для указания длины списка.

Это также можно сделать полиморфным по типу . В следующем примере входными данными является список тензоров (с длиной "N" ) одного (но неуказанного) типа ( "T" ), а выходными данными является один тензор соответствующего типа:

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

По умолчанию тензорные списки имеют минимальную длину 1. Вы можете изменить это значение по умолчанию, используя ограничение ">=" для соответствующего атрибута . В следующем примере входными данными является список как минимум из двух тензоров int32 :

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

Тот же синтаксис работает с атрибутами "list(type)" :

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

Входы и выходы

Подводя итог вышеизложенному, можно сказать, что регистрация операции может иметь несколько входов и выходов:

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

Каждая спецификация ввода или вывода имеет форму:

<name>: <io-type-expr>

где <name> начинается с буквы и может состоять из буквенно-цифровых символов и знаков подчеркивания. <io-type-expr> является одним из следующих выражений типа:

  • <type> , где <type> — поддерживаемый тип ввода (например, float , int32 , string ). Это указывает один тензор данного типа.

    См. tf.DType .

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

  • <attr-type> , где <attr-type> — это имя Attr с типом type или list(type) (с возможным ограничением типа). Этот синтаксис допускает полиморфные операции .

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

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

    Ссылка на атрибут типа list(type) позволяет вам принимать последовательность тензоров.

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

    Обратите внимание, что количество и типы тензоров на выходе out такие же, как и на входе in , поскольку оба имеют тип T .

  • Для последовательности тензоров одного типа: <number> * <type> , где <number> — это имя Attr с типом int . <type> может быть либо tf.DType , либо именем attr с типом type . В качестве примера первого, эта операция принимает список тензоров int32 :

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

    Принимая во внимание, что эта операция принимает список тензоров любого типа, если они все одинаковы:

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

  • Для ссылки на тензор: Ref(<type>) , где <type> — один из предыдущих типов.

Будет выведен любой атрибут, используемый в типе ввода. По соглашению эти предполагаемые атрибуты используют заглавные имена (например T или N ). В противном случае входы, выходы и атрибуты имеют имена, подобные параметрам функции (например, num_outputs ). Дополнительные сведения см. в предыдущем разделе об именовании .

Дополнительные сведения см. в файле tensorflow/core/framework/op_def_builder.h .

Обратная совместимость

Предположим, вы написали красивую индивидуальную операцию и поделились ею с другими, поэтому у вас есть довольные клиенты, использующие вашу операцию. Однако вы хотели бы каким-то образом внести изменения в операцию.

Как правило, изменения в существующих зарегистрированных спецификациях должны быть обратно совместимыми: изменение спецификации операции не должно нарушать предыдущие сериализованные буферы протокола GraphDef , созданные из более старых спецификаций. Подробности совместимости GraphDef описаны здесь .

Есть несколько способов сохранить обратную совместимость.

  1. Любые новые атрибуты, добавленные к операции, должны иметь определенные значения по умолчанию, и с этим значением по умолчанию операция должна иметь исходное поведение. Чтобы изменить операцию с неполиморфной на полиморфную, вы должны указать значение по умолчанию для нового типа attr, чтобы по умолчанию сохранить исходную подпись. Например, если ваша операция была:

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

    вы можете сделать его полиморфным обратно совместимым способом, используя:

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

  2. Вы можете безопасно сделать ограничение на attr менее строгим. Например, вы можете изменить {int32, int64} на {int32, int64, float} или type . Или вы можете заменить {"apple", "orange"} на {"apple", "banana", "orange"} или string .

  3. Вы можете изменить отдельные входы/выходы на входы/выходы списка, если значение по умолчанию для типа списка соответствует старой подписи.

  4. Вы можете добавить новый список ввода/вывода, если он по умолчанию пуст.

  5. Пространство имен любых новых операций, которые вы создаете, добавляя к именам операций префикс, уникальный для вашего проекта. Это позволяет избежать столкновения вашей операции с любыми операциями, которые могут быть включены в будущие версии TensorFlow.

  6. Планируйте заранее! Постарайтесь предвидеть будущее использование op. Некоторые изменения подписи не могут быть выполнены совместимым образом (например, преобразование списка одного типа в список различных типов).

Полный список безопасных и небезопасных изменений можно найти в tensorflow/core/framework/op_compatibility_test.cc . Если вы не можете сделать изменение операции обратно совместимым, создайте новую операцию с новым именем и новой семантикой.

Также обратите внимание, что хотя эти изменения могут поддерживать совместимость GraphDef , сгенерированный код Python может измениться таким образом, что он будет несовместим со старыми вызывающими программами. API Python можно сохранить совместимым путем тщательных изменений в написанной от руки оболочке Python, сохранив старую подпись, за исключением, возможно, добавления новых необязательных аргументов в конец. Как правило, несовместимые изменения могут быть сделаны только тогда, когда TensorFlow меняет основные версии, и они должны соответствовать семантике версий GraphDef .

Поддержка графического процессора

Вы можете реализовать разные ядра OpKernel и зарегистрировать одно для CPU, а другое для GPU, точно так же, как вы можете регистрировать ядра для разных типов . В tensorflow/core/kernels/ есть несколько примеров ядер с поддержкой GPU. Обратите внимание, что некоторые ядра имеют версию CPU в файле .cc , версию GPU в файле, оканчивающемся на _gpu.cu.cc , и некоторый общий код в файле .h .

Например, в tf.pad есть все, кроме ядра графического процессора, в tensorflow/core/kernels/pad_op.cc . Ядро графического процессора находится в tensorflow/core/kernels/pad_op_gpu.cu.cc , а общий код представляет собой шаблонный класс, определенный в tensorflow/core/kernels/pad_op.h . Мы организуем код таким образом по двум причинам: это позволяет использовать общий код между реализациями CPU и GPU, а реализация GPU помещается в отдельный файл, чтобы ее мог скомпилировать только компилятор GPU.

Следует отметить одну вещь: даже когда используется версия pad для ядра графического процессора, ей по-прежнему требуется ввод "paddings" в память ЦП. Чтобы отметить, что входы или выходы хранятся на ЦП, добавьте вызов HostMemory() в регистрацию ядра, например:

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

Компиляция ядра для устройства GPU

Посмотрите на cuda_op_kernel.cu.cc пример, использующий ядро ​​CUDA для реализации операции. tf_custom_op_library принимает аргумент gpu_srcs , в котором можно указать список исходных файлов, содержащих ядра CUDA (файлы *.cu.cc ). Для использования с двоичной установкой TensorFlow ядра CUDA должны быть скомпилированы с помощью компилятора NVIDIA nvcc . Вот последовательность команд, которые вы можете использовать для компиляции cuda_op_kernel.cu.cc и cuda_op_kernel.cc в единую динамически загружаемую библиотеку:

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 , созданный выше, может быть загружен как обычно в Python с помощью функции tf.load_op_library .

Обратите внимание: если ваши библиотеки CUDA не установлены в /usr/local/lib64 , вам нужно будет явно указать путь во второй команде (g++) выше. Например, добавьте -L /usr/local/cuda-8.0/lib64/ , если ваша CUDA установлена ​​в /usr/local/cuda-8.0 .

Реализовать градиент в Python

Имея граф операций, TensorFlow использует автоматическое дифференцирование (обратное распространение) для добавления новых операций, представляющих градиенты по отношению к существующим операциям. Чтобы автоматическая дифференциация работала для новых операций, вы должны зарегистрировать функцию градиента, которая вычисляет градиенты по отношению к входным данным операций по заданным градиентам по отношению к выходным данным операций.

Математически, если оператор вычисляет \(y = f(x)\) зарегистрированный градиент op преобразует градиенты \(\partial L/ \partial y\) потерь \(L\) относительно\(y\) в градиенты \(\partial L/ \partial x\) относительно \(x\) с помощью цепного правила:

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

В случае ZeroOut только одна запись на входе влияет на выход, поэтому градиент по отношению к входу представляет собой разреженный «горячий» тензор. Это выражается следующим образом:

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

Подробности о регистрации функций градиента с помощью tf.RegisterGradient :

  • Для оператора с одним выходом функция градиента примет tf.Operation , op и tf.Tensor grad и построит новые операции из тензоров op.inputs[i] , op.outputs[i] и grad . Информацию о любых атрибутах можно найти через tf.Operation.get_attr .

  • Если у операции есть несколько выходов, функция градиента примет op и grads , где grads — это список градиентов по отношению к каждому выходу. Результатом функции градиента должен быть список объектов Tensor , представляющих градиенты по отношению к каждому входу.

  • Если для некоторых входных данных нет четко определенного градиента, например для целочисленных входных данных, используемых в качестве индексов, соответствующий возвращаемый градиент должен быть None . Например, для операции, принимающей тензор с плавающей запятой x и целочисленный индекс i , функция градиента return [x_grad, None] .

  • Если для оператора вообще не существует значимого градиента, вам часто не придется регистрировать какой-либо градиент, и пока градиент оператора никогда не понадобится, все будет в порядке. В некоторых случаях оператор не имеет четко определенного градиента, но может участвовать в вычислении градиента. Здесь вы можете использовать ops.NotDifferentiable для автоматического распространения нулей в обратном направлении.

Обратите внимание, что во время вызова функции градиента доступен только граф потока данных операций, а не сами данные тензора. Таким образом, все вычисления должны выполняться с использованием других операций тензорного потока, которые должны выполняться во время выполнения графа.

Функции формы в C++

API TensorFlow имеет функцию, называемую «вывод формы», которая предоставляет информацию о формах тензоров без необходимости выполнения графика. Вывод формы поддерживается «функциями формы», которые регистрируются для каждого типа операции в объявлении C++ REGISTER_OP и выполняют две роли: утверждают, что формы входных данных совместимы во время построения графа, и определяют формы для выходных данных.

Функции формы определяются как операции над классом shape_inference::InferenceContext . Например, в функции shape для ZeroOut:

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

c->set_output(0, c->input(0)); объявляет, что форма первого вывода должна быть установлена ​​на форму первого входа. Если вывод выбирается по его индексу, как в приведенном выше примере, вторым параметром set_output должен быть объект ShapeHandle . Вы можете создать пустой объект ShapeHandle с помощью его конструктора по умолчанию. Объект ShapeHandle для ввода с индексом idx можно получить c->input(idx) .

Существует ряд общих функций формы, которые применяются ко многим операциям, например shape_inference::UnchangedShape , которую можно найти в common_shape_fns.h и использовать следующим образом:

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

Функция формы также может ограничивать форму ввода. Для версии ZeroOut с ограничением формы вектора функция формы будет следующей:

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

Вызов WithRank проверяет, что входная фигура c->input(0) имеет фигуру ровно с одним измерением (или, если входная фигура неизвестна, выходная фигура будет вектором с одним неизвестным измерением).

Если ваша операция является полиморфной с несколькими входными данными , вы можете использовать члены InferenceContext , чтобы определить количество проверяемых фигур, и Merge , чтобы проверить, что все фигуры совместимы (в качестве альтернативы, доступ к атрибутам, указывающим длину, с помощью InferenceContext::GetAttr , который обеспечивает доступ к атрибутам операции).

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

Поскольку вывод формы является дополнительной функцией, а формы тензоров могут динамически изменяться, функции формы должны быть устойчивыми к неполной информации о форме для любого из входных данных. Метод Merge в InferenceContext позволяет вызывающему объекту утверждать, что две фигуры одинаковы, даже если одна или обе из них не имеют полной информации. Функции формы определены для всех основных операций TensorFlow и предоставляют множество различных примеров использования.

Класс InferenceContext имеет ряд функций, которые можно использовать для определения манипуляций с функциями формы. Например, вы можете проверить, что конкретное измерение имеет очень конкретное значение, используя InferenceContext::Dim и InferenceContext::WithValue ; вы можете указать, что выходное измерение является суммой/произведением двух входных измерений, используя InferenceContext::Add и InferenceContext::Multiply . См. класс InferenceContext для всех различных манипуляций с фигурами, которые вы можете указать. В следующем примере форма первого вывода устанавливается равной (n, 3), где первый вход имеет форму (n, ...)

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

Если у вас есть сложная функция формы, вам следует подумать о добавлении теста для проверки того, что различные комбинации входных форм создают ожидаемые комбинации выходных форм. Вы можете увидеть примеры написания этих тестов в некоторых наших тестах core ops . (Синтаксис INFER_OK и INFER_ERROR немного загадочен, но постарайтесь быть компактным при представлении входных и выходных спецификаций формы в тестах. Пока посмотрите окружающие комментарии в этих тестах, чтобы получить представление о спецификации строки формы).

Создайте пакет pip для своей пользовательской операции

Чтобы создать пакет pip для своей операции, см. пример tensorflow/custom-op . В этом руководстве показано, как создавать собственные операции из пакета pip TensorFlow вместо сборки TensorFlow из исходного кода.