【架构分析】TensorFlow 自定义算子实现原理分析

目录

背景介绍

代码示例

核心数据结构

REGISTER_OP 分析

REGISTER_KERNEL_BUILDER 分析

自定义算子运行分析

---

 

背景介绍

TensorFlow官网如何创建自定义算子OP的How to文档链接，本文基于该文档中的代码示例，着重分析TensorFlow框架是如何实现用户自定义算子扩展的功能，力求知其然还要知其所以然。

TensorFlow源码下载链接，本文的分析基于最新的master分支 2.3.2 版本 (Google 自己又整了个bazel编译框架，编译过程也挺折腾的... ）

 

代码示例

//TfCustomOp.cpp

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;


//Define&Register OP
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();
    });


//Define OP's Kernel
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();

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

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

//Register OP's Kernel
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

#CMakeLists.txt

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(tfcustomop_project)

# Define our library target
add_library(tfcustomop SHARED TfCustomOp.cpp print_stack.cpp)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -D_GLIBCXX_USE_CXX11_ABI=1")

include_directories(/tensorflow/include)

# Link against LibTorch
#target_link_libraries(tfcustomop "/tensorflow/libtensorflow_framework.so.2")

从示例代码可以看到创建自定义算子OP有三步

1. 使用REGISTER_OP宏定义OP算子的名字和属性，注册到TensorFlow框架
2. 从OpKernel类继承创建自己的Kernel计算子类，将应用的计算业务逻辑实现到Compute函数
3. 使用REGISTER_KERNEL_BUILDER宏注册自定义的Kernel名字和属性，注册到TensorFlow框架

这种用宏扩展自定义OP的方式非常类似于我上一篇文章中介绍的pytorch TORCH_LIBRARY 宏，实际上经过对TensorFlow框架代码的分析也确认两者的设计思想很接近（Google和Facebook里面都是编程高手 -v- )；区别在于TensorFlow的设计中，算子和Kernel计算是两个分开的概念，其中算子的语义是平台无关的，而算子对应的Kernel计算则是跟平台（CPU、GPU、TPU）相关的。这样，就可以很方便的对语义不变的算子提供不同平台的具体实现了

经过cmake编译后生成了libtfcustomop.so，使用下面的python测试程序验证结果确认自定义算子和Kernel已经生效

import tensorflow as tf

mylib = tf.load_op_library("path/to/libtfcustomop.so")

with tf.compat.v1.Session():
    print(mylib.zero_out([4,4,4,4,4]).eval())

[4,0,0,0,0]

 

核心数据结构

算子与Kernel定义与注册核心数据结构

在详细结果算子与Kernel的定义与注册原理前，概况相关核心的数据结构和概念如下

• OpXXX的数据结构负责构建自定义算子(OpDef)，最终注册自定义算子的描述对象到OpRegistry中的集合便于后续查找使用
• OpKernelXXX的数据结构负责构建自定义Kernel(OpKernel)，最终注册自定义Kernel的描述对象到KernelRegistry中的集合便于后续查找
• TensorFlow框架代码中存在几个核心的Def数据结果：GraphDef NodeDef OpDef OpKernel；GraphDef和OpKernel顾名思义容易理解，其中NodeDef和OpDef的含义和差别在于： 在TensorFlow中，任何一个操作都有两部分组成。一个是OpDef类型的对象，它类似于C语言的声明用来实现操作的业务逻辑，通过OpDef对象，我们可以告诉TensorFlow当前操作的输入、输出以及参数的类型。NodeDef对象则是用来控制输入与输出的传递，在Nodef对象中，会指定当前操作的输入是由哪些操作的输出给定；当前操作是输出，又将作为哪些操作的输入。

有了上面的基本知识，下面逐步分析TensorFlow框架中的实现原理

 

REGISTER_OP 分析

REGISTER_OP宏负责定义与注册用户的OP算子，它的定义在/tensorflow/core/framework/op.h，通过gcc工具展开宏定义

g++ -I /tensorflow/include -std=c++11 -E TfCustomOp.cpp

static ::tensorflow::register_op::OpDefBuilderReceiver register_op0 __attribute__((unused)) = ::tensorflow::register_op::OpDefBuilderWrapper("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

可以看到REGISTER_OP宏通过构建了static OpDefBuilderReceiver变量，实现了OP的定义与注册过程，整个过程的时序图如下

自定义OP注册时序

 

REGISTER_KERNEL_BUILDER 分析

REGISTER_KERNEL_BUILDER宏负责自定义Kernel的注册，它的定义在/tensorflow/core/framework/op_kernel.h，通过gcc工具展开宏定义

constexpr bool should_register_1__flag = true; 
static ::tensorflow::kernel_factory::OpKernelRegistrar registrar__body__1__object( should_register_1__flag ? ::tensorflow::register_kernel::Name("ZeroOut").Device(DEVICE_CPU).Build() : nullptr, 
"ZeroOutOp", 
[](::tensorflow::OpKernelConstruction* context) -> ::tensorflow::OpKernel* 
{ return new ZeroOutOp(context); });;

通过构建了static OpKernelRegistrar变量，实现了Kernel的注册过程，整个过程的时序图如下

自定义Kernel注册时序

 

自定义算子运行分析

经过自定义OP的注册，Kernel注册，Kernel Compute实现业务逻辑计算这三步后，就可以在python程序中使用自定义的OP算子了，核心的时序如下图

自定义算子查找与运行时序