Почему и как разработчики расширяют TensorFlow с помощью пользовательских ядер
data-science tensorflow tensorflow-custom-op tensorflow-gpu-kernel tensorflow-c++-ops tensorflow-cpu-kernel median-pooling-op bazel-tensorflow-compile tf.load_op_library

Почему и как разработчики расширяют TensorFlow с помощью пользовательских ядер

25 июля 2025 г.

Обзор контента

  • Предварительные условия
  • Определите интерфейс OP
  • Внедрить ядро для OP
  • Многопоточные ядра процессора
  • Ядра графического процессора
  • Создать библиотеку OP
  • Скомпилируйте OP, используя системный компилятор (бинарная установка TensorFlow)
  • Скомпилируйте OP с использованием Bazel (установка источника TensorFlow)
  • Используйте OP в Python
  • Убедитесь, что OP работает

Примечание:Чтобы гарантировать, что ваши пользовательские Ops C ++ совместим с официальными PIP -пактами TensorFlow, пожалуйста, следуйте руководству вПользовательский репозиторий OPПолем Он имеет пример кодового кода, а также изображения Docker для построения и распространения ваших пользовательских OPS.

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

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

Например, представьте, что вы хотите реализовать что -то вроде «медианного объединения», похожего на оператора «maxpool», но вычисляют медианы над раздвижными окнами вместо максимальных значений. Это может быть возможно с использованием композиции операций (например, с использованием ExtractImagePatches и Topk), но может быть не таким эффективным или памятью, как нативная операция, где вы можете сделать что-то более умное в одной, плавной операции. Как всегда, обычно сначала стоит попытаться выразить то, что вы хотите, используя композицию оператора, только выбор добавить новую операцию, если это оказалось трудным или неэффективным.

Чтобы включить свой пользовательский OP, вам понадобится:

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

Предварительные условия

  • Некоторое знакомство с C ++.
  • Должно быть, установилTensorflow Binary, или должен иметьЗагруженный источник TensorFlowи иметь возможность построить его.

Определите интерфейс OP

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

Чтобы увидеть, как это работает, предположим, что вы хотели бы создать OP, который берет тензорint32S и выводит копию тензора, со всеми, кроме первого элемента, установленного на ноль. Для этого создайте файл с именемzero_out.ccПолем Затем добавьте звонок вREGISTER_OPМакро, который определяет интерфейс для вашего 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_zero32-битных целых чисел в качестве входных данных и выводит тензорzeroed32-битных целых чисел. OP также использует функцию формы, чтобы гарантировать, что выходной тензор является той же формой, что и входной тензор. Например, если вход является тензором формы [10, 20], то эта функция формы указывает, что выходная форма также является [10, 20].

Примечание:Имя OP должно быть в Camelcase, и оно должно быть уникальным среди всех других OP, которые зарегистрированы в бинарном.

Внедрить ядро для OP

После определения интерфейса предоставьте одну или несколько реализаций 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. В регистрации вы указываете различные ограничения, при которых будет работать это ядро. Например, у вас может быть одно ядро, сделанное для процессоров, и отдельное для графических процессоров.

Сделать это дляZeroOutop, добавьте следующее вzero_out.cc:

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

Важный:Экземпляры вашего Opkernel можно получить одновременно. ТвойComputeМетод должен быть защищен потоком. Охраняйте любой доступ к классным членам с помощью Mutex. Или еще лучше, не делитесь штатом через классных членов! Подумайте об использованииResourceMgrотслеживать состояние OP.

Многопоточные ядра процессора

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

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

Ядро графического процессора реализовано в двух частях: Opkernel и Cuda ядро и его код запуска.

Иногда реализация Opkernel является общей между процессором и ядром графического процессора, например, вокруг проверки входов и выделения выходов. В этом случае предложенная реализация - это:

  1. Определите шаблон Opkernel на устройстве и примитивный тип тензора.
  2. Чтобы выполнить фактическое вычисление выхода, функция вычисления вызывает структуру шаблона.
  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

Создать библиотеку OP

Скомпилируйте OP, используя системный компилятор (бинарная установка TensorFlow)

Вы должны быть в состоянии компилироватьzero_out.ccсC++компилятор, такой какg++илиclangДоступно в вашей системе. Бинарный пакет PIP устанавливает файлы заголовков и библиотеку, которая необходимо для составления вашего OP в местах, которые являются специфичными для системы. Тем не менее, библиотека 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++Установлен, вот последовательность команд, которые вы можете использовать для составления вашего OP в динамическую библиотеку.

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 при строительстве требуется дополнительный флаг "-nuded Dynamic_lookup".soфайл.

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

Скомпилируйте OP с использованием Bazel (установка источника TensorFlow)

Если у вас установлены источники TensorFlow, вы можете использовать систему сборки TensorFlow для составления вашего OP. Поместите файл сборки с следующим правилом сборки 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Полем Поместите файл сборки со следующим правилом сборки 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

Примечание:Как объяснено выше, если вы компилируете с GCC> = 5 добавить--cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0"к аргументам командной строки Bazel.

Примечание:Хотя вы можете создать общую библиотеку (.soфайл) со стандартомcc_libraryправило, мы настоятельно рекомендуем вам использоватьtf_custom_op_libraryмакро. Он добавляет некоторые необходимые зависимости и выполняет проверки, чтобы убедиться, что общая библиотека совместима с механизмом загрузки плагина TensorFlow.

Используйте OP в Python

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

Чтобы сделать OP доступным в качестве обычной функции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

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

Хороший способ убедиться, что вы успешно реализовали свой OP - это написать тест для него. Создайте файл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

Первоначально опубликовано наTensorflowВеб -сайт, эта статья появляется здесь под новым заголовком и имеет лицензию в CC на 4.0. Образцы кода, разделенные по лицензии Apache 2.0.

Оригинал

Recent Post

Categories

PREVIOUS ARTICLE
NEXT ARTICLE