Если вы хотите создать операцию, которая не поддерживается существующей библиотекой TensorFlow, мы рекомендуем сначала попробовать написать операцию на Python как композицию существующих операций или функций Python. Если это невозможно, вы можете создать собственную операцию C++. Существует несколько причин, по которым вам может потребоваться создать собственную операцию C++:
- Непросто и невозможно выразить вашу операцию как композицию существующих операций.
- Неэффективно выражать операцию как композицию существующих примитивов.
- Вы хотите вручную объединить композицию примитивов, которую будущему компилятору будет сложно объединить.
Например, представьте, что вы хотите реализовать что-то вроде «объединения медиан», аналогичного оператору «MaxPool», но вычислять медианы по скользящим окнам, а не по максимальным значениям. Выполнение этого с использованием композиции операций может быть возможным (например, с использованием ExtractImagePatches и TopK), но оно может быть не таким эффективным с точки зрения производительности или памяти, как собственная операция, где вы можете сделать что-то более умное в одной объединенной операции. Как всегда, обычно сначала стоит попытаться выразить то, что вы хотите, с помощью композиции операторов, добавляя новую операцию только в том случае, если это окажется сложным или неэффективным.
Чтобы включить свою собственную операцию, вам необходимо:
- Зарегистрируйте новую операцию в файле C++. Регистрация операции определяет интерфейс (спецификацию) для функциональности операции, которая не зависит от реализации операции. Например, регистрация операции определяет имя операции, а также ее входы и выходы. Он также определяет функцию формы, которая используется для вывода формы тензора.
- Реализуйте операцию на C++. Реализация операции называется ядром и представляет собой конкретную реализацию спецификации, которую вы зарегистрировали на шаге 1. Может существовать несколько ядер для разных типов ввода/вывода или архитектур (например, ЦП, ГП).
- Создайте оболочку Python (необязательно). Эта оболочка представляет собой общедоступный API, который используется для создания операции в Python. На основе регистрации операции создается оболочка по умолчанию, которую можно использовать напрямую или добавлять.
- Напишите функцию для вычисления градиентов для операции (необязательно).
- Проверьте оп. Обычно для удобства мы делаем это на Python, но вы также можете протестировать эту операцию на C++. Если вы определяете градиенты, вы можете проверить их с помощью Python
tf.test.compute_gradient_error
. См.relu_op_test.py
как пример проверки прямых функций Relu-подобных операторов и их градиентов.
Предварительные условия
- Некоторое знакомство с C++.
- Должен быть установлен двоичный файл TensorFlow или загружен исходный код TensorFlow и иметь возможность его собрать.
Определить операционный интерфейс
Вы определяете интерфейс операции, регистрируя ее в системе 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].
Реализовать ядро для операции
После определения интерфейса предоставьте одну или несколько реализаций операции. Чтобы создать одно из этих ядер, создайте класс, расширяющий 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
. Эта функция распределяет вычислительную функцию по потокам, сконфигурированным для использования внутриоперационной многопоточности (см.tra_op_parallelism_threads в config.proto
).
Ядра графического процессора
Ядро графического процессора состоит из двух частей: OpKernel, ядра CUDA и его кода запуска.
Иногда реализация OpKernel является общей для ядра ЦП и графического процессора, например, при проверке входных данных и распределении выходных данных. В этом случае предлагаемая реализация заключается в следующем:
- Определите шаблон OpKernel на устройстве и примитивный тип тензора.
- Чтобы выполнить фактическое вычисление выходных данных, функция Compute вызывает шаблонную структуру-функтор.
- Специализация этого функтора для 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 и пытаетесь скомпилировать свою библиотеку op с помощью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
API TensorFlow Python предоставляет функцию 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()
Затем запустите тест (при условии, что у вас установлен tensorflow):
$ python zero_out_op_test.py
Встройте расширенные функции в свою работу
Теперь, когда вы знаете, как построить базовую (и несколько ограниченную) операцию и ее реализацию, мы рассмотрим некоторые из более сложных вещей, которые вам обычно придется встроить в свою операцию. Это включает в себя:
- Условные проверки и валидация
- Операционная регистрация
- Поддержка графического процессора
- Реализация градиента в Python
- Функции формы в C++
Условные проверки и валидация
В приведенном выше примере предполагалось, что операция op применима к тензору любой формы. Что, если это применимо только к векторам? Это означает добавление проверки в приведенную выше реализацию 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_);
}
Типы атрибутов
В атрибуте поддерживаются следующие типы:
-
string
: любая последовательность байтов (не обязательно UTF8). -
int
: целое число со знаком. -
float
: число с плавающей запятой. -
bool
: Истина или ложь. -
type
: одно из (не ссылочных) значенийDataType
. -
shape
:TensorShapeProto
. -
list(<type>)
: список<type>
, где<type>
— один из вышеуказанных типов. Обратите внимание, чтоlist(list(<type>))
недействителен.
См. также: op_def_builder.cc:FinalizeAttr
для получения окончательного списка.
Значения и ограничения по умолчанию
Атрибуты могут иметь значения по умолчанию, а некоторые типы атрибутов могут иметь ограничения. Чтобы определить атрибут с ограничениями, вы можете использовать следующие <attr-type-expr>
:
{'<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
. В этом примере атрибут 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>
. Например, следующая регистрация операции указывает, что 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
Именование
Входные, выходные данные и атрибуты обычно должны иметь имена Snake_case. Единственным исключением являются атрибуты, которые используются в качестве типа ввода или типа вывода. Эти атрибуты можно вывести, когда операция добавляется в график, и поэтому они не отображаются в функции операции. Например, это последнее определение 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.
Сравните это с операцией, у которой есть тип attr, определяющий тип вывода:
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
Перечислить входы и выходы
Помимо возможности принимать или создавать разные типы, операции могут потреблять или создавать переменное количество тензоров.
В следующем примере атрибут 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
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
, либо именем атрибута с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
описаны здесь .
Есть несколько способов сохранить обратную совместимость.
Любые новые атрибуты, добавленные в операцию, должны иметь определенные значения по умолчанию, и с этим значением по умолчанию операция должна иметь исходное поведение. Чтобы изменить операцию с неполиморфной на полиморфную, вы должны задать значение по умолчанию для нового типа attr, чтобы сохранить исходную подпись по умолчанию. Например, если ваша операция была:
REGISTER_OP("MyGeneralUnaryOp") .Input("in: float") .Output("out: float");
вы можете сделать его полиморфным обратно совместимым способом, используя:
REGISTER_OP("MyGeneralUnaryOp") .Input("in: T") .Output("out: T") .Attr("T: numerictype = DT_FLOAT");
Вы можете безопасно сделать ограничение атрибута менее строгим. Например, вы можете изменить
{int32, int64}
на{int32, int64, float}
илиtype
. Или вы можете изменить значение с{"apple", "orange"}
на{"apple", "banana", "orange"}
илиstring
.Вы можете изменить отдельные входы/выходы на входы/выходы списка, если значение по умолчанию для типа списка соответствует старой сигнатуре.
Вы можете добавить новый список ввода/вывода, если он по умолчанию пуст.
Пространство имен для любых новых операций, которые вы создаете, добавляя к именам операций что-то уникальное для вашего проекта. Это позволяет избежать конфликта вашей операции с любыми операциями, которые могут быть включены в будущие версии TensorFlow.
Планируйте заранее! Постарайтесь предвидеть будущее использование операции. Некоторые изменения сигнатур невозможно выполнить совместимым способом (например, превратить список одного типа в список разных типов).
Полный список безопасных и небезопасных изменений можно найти в tensorflow/core/framework/op_compatibility_test.cc
. Если вы не можете сделать изменения в операции обратно совместимыми, создайте новую операцию с новым именем и новой семантикой.
Также обратите внимание, что хотя эти изменения могут обеспечить совместимость GraphDef
, сгенерированный код Python может измениться таким образом, что это будет несовместимо со старыми вызывающими объектами. Совместимость API Python можно сохранить путем осторожных изменений в написанной вручную оболочке Python, сохранив старую подпись, за исключением, возможно, добавления в конец новых необязательных аргументов. Обычно несовместимые изменения могут быть внесены только тогда, когда TensorFlow меняет основные версии, и они должны соответствовать семантике версии GraphDef
.
Поддержка графического процессора
Вы можете реализовать разные OpKernels и зарегистрировать один для CPU, а другой для GPU, точно так же, как вы можете зарегистрировать ядра для разных типов . В tensorflow/core/kernels/
есть несколько примеров ядер с поддержкой графического процессора. Обратите внимание, что некоторые ядра имеют версию ЦП в файле .cc
, версию графического процессора в файле, заканчивающемся на _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
. Мы организуем код таким образом по двум причинам: он позволяет использовать общий код для реализаций ЦП и графического процессора, а также помещает реализацию графического процессора в отдельный файл, чтобы его можно было скомпилировать только компилятором графического процессора.
Следует отметить одну вещь: даже когда используется версия pad
для ядра графического процессора, ей все равно требуется ввод "paddings"
в память ЦП. Чтобы отметить, что входные или выходные данные хранятся в ЦП, добавьте вызов HostMemory()
к регистрации ядра, например:
#define REGISTER_GPU_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("Pad") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T") \
.HostMemory("paddings"), \
PadOp<GPUDevice, T>)
Компиляция ядра для устройства с графическим процессором
Посмотрите 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)\) зарегистрированная операция градиента преобразует градиенты \(\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 имеет несколько выходов, функция градиента будет принимать
op
иgrads
, гдеgrads
— это список градиентов по отношению к каждому выходу. Результатом функции градиента должен быть список объектовTensor
, представляющих градиенты по отношению к каждому входу.Если для некоторых входных данных нет четко определенного градиента, например для целочисленных входных данных, используемых в качестве индексов, соответствующий возвращаемый градиент должен быть
None
. Например, для операции, принимающей тензор с плавающей запятойx
и целочисленный индексi
, функция градиентаreturn [x_grad, None]
.Если для операции вообще нет значимого градиента, вам часто не придется регистрировать какой-либо градиент, и пока градиент операции никогда не понадобится, все будет в порядке. В некоторых случаях операция не имеет четко определенного градиента, но может участвовать в вычислении градиента. Здесь вы можете использовать
ops.NotDifferentiable
для автоматического распространения нулей в обратном направлении.
Обратите внимание, что во время вызова функции градиента доступен только граф потока данных операций, а не сами данные тензора. Таким образом, все вычисления должны выполняться с использованием других операций тензорного потока, которые будут выполняться во время выполнения графа.
Добавляйте подсказки типов при регистрации пользовательского градиента для типа операции, чтобы сделать код более читабельным, отлаживаемым, простым в обслуживании и более надежным благодаря проверке данных. Например, при использовании op
в качестве параметра функции укажите, что функция градиента будет принимать tf.Operation
в качестве типа параметра.
Функции формы в C++
API TensorFlow имеет функцию под названием «вывод формы», которая предоставляет информацию о формах тензоров без необходимости выполнения графа. Вывод формы поддерживается «функциями формы», которые зарегистрированы для каждого типа операции в объявлении C++ REGISTER_OP
и выполняют две роли: утверждение совместимости входных форм во время построения графа и указание форм для выходных данных.
Функции формы определяются как операции в классе shape_inference::InferenceContext
. Например, в функции формы для 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();
});
Если у вас сложная функция формы, вам следует рассмотреть возможность добавления теста для проверки того, что различные комбинации входных форм создают ожидаемые комбинации выходных форм. Вы можете увидеть примеры написания этих тестов в некоторых наших основных операционных тестах . (Синтаксис INFER_OK
и INFER_ERROR
немного загадочный, но постарайтесь быть компактным при представлении спецификаций входных и выходных форм в тестах. А пока посмотрите комментарии в этих тестах, чтобы получить представление о спецификации строки формы).
Создайте пакет pip для вашей пользовательской операции.
Чтобы создать пакет pip
для вашей операции, см. пример tensorflow/custom-op . В этом руководстве показано, как создавать собственные операции из пакета pip TensorFlow вместо сборки TensorFlow из исходного кода.