本文作者:李杰
TF计算图从逻辑层来讲,由op与tensor构成。op是项点代表计算单元,tensor是边代表op之间流动的数据内容,两者配合以数据流图的形式来表达计算图。那么op对应的物理层实现是什么?TF中有哪些op,以及各自的适用场景是什么?op到底是如何运行的?接下来让我们一起探索和回答这些问题。
一、初识op
1.1 op定义
op代表计算图中的节点,是tf.Operation对象,代表一个计算单元。用户在创建模型和训练代码时,会创建一系列op及其依赖关系,并将这些op和依赖添加到tf.Graph对象中(一般为默认图)。比如:tf.matmul()就是一个op,它有两个输入tensor和一个输出tensor。
1.2 op分类
op的分类一般有多个视角,比如按是否内置划分、按工作类型划分。
按是否内置划分,一般分为:内置op和自定义op(见“二、自定义op”部分介绍)。
按工作类型划分,一般分为:常见数学op、数组op、矩阵op、有状态op、神经网络op、检查点op、队列与同步op、控制流op。TF白皮书对内置op的分类总结如下:
1.3 op与kernel
op一般都有名称且代表一个抽象的计算过程。op可以设置若干属性,但这些属性必须在编译期提供或推理得到,因为它们用来实例化一个节点对象从而执行真正的计算。属性的经典用法就是拿来支持类型多态,比如两个浮点张量的矩阵乘法与两个整型张量的矩阵乘法。
kernel是op在指定设备类型(CPU/GPU)上的具体实现。TF二进制库通过注册机制定义了一系列op及对应的kernel实现,用户可以提供额外的op定义与kernel实现进行扩充。一般来说,一个op对应多个kernel实现。
接下来让我们一起用矩阵乘法MatMul算子的相关代码来理解op与kernel的关系(此处不必纠结代码细节,只需体会op与kernel关系即可):
// 首先给出op注册的定义。其中输入输出支持泛型,其合法类型在Attr中进行枚举。 // 代码位置 tensorflow1.15.5tensorflowcoreopsmath_ops.cc REGISTER_OP("MatMul") .Input("a: T") .Input("b: T") .Output("product: T") .Attr("transpose_a: bool = false") .Attr("transpose_b: bool = false") .Attr( "T: {bfloat16, half, float, double, int32, int64, complex64, " "complex128}") .SetShapeFn(shape_inference::MatMulShape); // MatMul的实现,采用类模板机制 // 代码位置 tensorflow1.15.5tensorflowcorekernelsmatmul_op.cc template <typename Device, typename T, bool USE_CUBLAS> class MatMulOp : public OpKernel { public: explicit MatMulOp(OpKernelConstruction* ctx) : OpKernel(ctx), algorithms_set_already_(false) { OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_)); LaunchMatMul<Device, T, USE_CUBLAS>::GetBlasGemmAlgorithm( ctx, &algorithms_, &algorithms_set_already_); use_autotune_ = MatmulAutotuneEnable(); } // 省略了很多代码... private: std::vector<int64> algorithms_; bool algorithms_set_already_; bool use_autotune_; bool transpose_a_; bool transpose_b_; }; // MatMul的op定义与kernel实现绑定处理 // 代码位置 tensorflow1.15.5tensorflowcorekernelsmatmul_op.cc #define REGISTER_CPU_EIGEN(T) /*cpu与eigen组合对应实现*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); #define REGISTER_CPU(T) /*cpu对应实现(eigen与非eigen)*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T"), MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); REGISTER_CPU_EIGEN(T); #define REGISTER_GPU(T) /*gpu对应实现(cublas与非cublas)*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_GPU).TypeConstraint<T>("T"), MatMulOp<GPUDevice, T, true /* cublas, true by default */>); REGISTER_KERNEL_BUILDER(Name("MatMul") .Device(DEVICE_GPU) .TypeConstraint<T>("T") .Label("cublas"), MatMulOp<GPUDevice, T, true /* cublas */>)
二、自定义op
用户编写的模型训练代码一般由TF原生的op算子及其依赖关系组成,但有时候我们定义的计算逻辑在TF中没有相应的op实现。根据TensorFlow官网的建议,我们应当先组合python op算子或python函数进行尝试。完成尝试之后再决定要不要自定义op。
2.1 自定义op场景
一般来说,需要自定义op的场景有如下3个:
在此举个例子方便大家理解。假如我们要实现一个新计算实逻:中位数池化(median pooling),过程中要在滑动窗口不断求得中位数。检索TF文档没有发现对应op,因此我们先考虑用TF python op组合来实现它,果然通过ExtractImagePatches and TopK就可以实现这个功能。经测试前述组合方案并不是计算和存储高效的,因此我们就有必要将median pooling在一个op中进行高效实现。
2.2 自定义op流程
自定义op一般遵循5个基本步骤:
接下来我们就以官网最简单的ZeroOut同步式自定义op(继承OpKernel)为例,结合代码来讲述上述5个步骤。下面先给出步骤1和步骤2用C++实现的代码(官方推荐用bazel编译so文件):
// 步骤1:注册op REGISTER_OP("ZeroOut") .Input("to_zero: int32") .Output("zeroed: int32") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); //c's input and output type is std::vector<ShapeHandle> return Status::OK(); }); // 步骤2:定义kernel(常规CPU设备),并把kernel与op绑定 class ZeroOutOp : public OpKernel { public: explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {} void Compute(OpKernelContext* context) override { // Grab the input tensor from OpKernelContext instance 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)); // OP_REQUIRES_OK第二个参数一般为方法调用,此处为输出张量分配内存空间 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); } }; REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);
步骤3加载上述so文件(自动完成前后端op映射);步骤4是可选项,此处不需要;步骤5基于python api测试op功能。相应代码如下:
import tensorflow as tf zero_out_module = tf.load_op_library('./zero_out.so') # 加载so文件生成python module with tf.Session(''): zero_out_module.zero_out([[1, 2], [3, 4]]).eval() # Prints array([[1, 0], [0, 0]], dtype=int32)
2.3 高级话题
关于op的技术话题还有很多,我们在此简述一些要点:
三、op工作原理
3.1 op运行框架
整体来看,op与kernel都有其结构描述与统一的注册管理中心。而OpDefBuilder有两个包装类OpDefBuilderWrapper和OpDefBuilderReceiver,前者支持op构建的链式语法,后者接受op构建结果并进行注册。众所周知,op是编译期概念,而kernel是运行期概念,在AI编译器的后端处理流程中会进行op的算子选择,此过程会基于一系列策略为op匹配最合适的kernel实现。
3.2 若干技术细节
首先,我们来看一下大家在使用TensorFlow过程中经常碰到的libtensorflow_framework.so。按照tf1.15.5/tensorflow/BUILD中的描述,libtensorflow_framework.so定义了op和kernel的注册机制而不涉及具体实现。
// rootdir=tensorflow1.15.5 // ${rootdir}/tensorflow/BUILD /* # A shared object which includes registration mechanisms for ops and # kernels. Does not include the implementations of any ops or kernels. Instead, # the library which loads libtensorflow_framework.so # (e.g. _pywrap_tensorflow_internal.so for Python, libtensorflow.so for the C # API) is responsible for registering ops with libtensorflow_framework.so. In # addition to this core set of ops, user libraries which are loaded (via # TF_LoadLibrary/tf.load_op_library) register their ops and kernels with this # shared object directly. */ tf_cc_shared_object( name = "tensorflow_framework", framework_so = [], linkopts = select({ "//tensorflow:macos": [], "//tensorflow:windows": [], "//tensorflow:freebsd": [ "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)", "-lexecinfo", ], "//conditions:default": [ "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)", ], }), linkstatic = 1, per_os_targets = True, soversion = VERSION, visibility = ["//visibility:public"], deps = [ "//tensorflow/cc/saved_model:loader_lite_impl", "//tensorflow/core:core_cpu_impl", "//tensorflow/core:framework_internal_impl", /* 展开此target进行查看 */ "//tensorflow/core:gpu_runtime_impl", "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry_impl", "//tensorflow/core:lib_internal_impl", "//tensorflow/stream_executor:stream_executor_impl", "//tensorflow:tf_framework_version_script.lds", ] + tf_additional_binary_deps(), ) // ${rootdir}/tensorflow/core/BUILD tf_cuda_library( name = "framework_internal_impl", srcs = FRAMEWORK_INTERNAL_PRIVATE_HEADERS + glob( // 可以查看FRAMEWORK_INTERNAL_PRIVATE_HEADERS内容 [ "example/**/*.cc", "framework/**/*.cc", "util/**/*.cc", "graph/edgeset.cc", "graph/graph.cc", "graph/graph_def_builder.cc", "graph/node_builder.cc", "graph/tensor_id.cc", "graph/while_context.h", "graph/while_context.cc", ], // 省略了诸多代码 ) // FRAMEWORK_INTERNAL_PRIVATE_HEADERS的内容 FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [ "graph/edgeset.h", "graph/graph.h", "graph/graph_def_builder.h", "graph/node_builder.h", "graph/tensor_id.h", ] + glob( [ "example/**/*.h", "framework/**/*.h", // 这里就是重点,查看${rootdir}/tensorflow/core/framework/op.h和opkernel.h "util/**/*.h", ] ) // 先来看op.h #define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name) #define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name) #define REGISTER_OP_UNIQ(ctr, name) static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr TF_ATTRIBUTE_UNUSED = ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( name)>(name) // 再来看看opkernel.h #define REGISTER_KERNEL_BUILDER(kernel_builder, ...) REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__) #define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__) #define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) constexpr bool should_register_##ctr##__flag = SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__); static ::tensorflow::kernel_factory::OpKernelRegistrar registrar__body__##ctr##__object( should_register_##ctr##__flag ? ::tensorflow::register_kernel::kernel_builder.Build() : nullptr, #__VA_ARGS__, [](::tensorflow::OpKernelConstruction* context) -> ::tensorflow::OpKernel* { return new __VA_ARGS__(context); }); 参照上述同样的流程,我们可以发现libtensorflow.so中涉及op与kernel的具体实现,同时也包括Session的具体实现。 最后,我们再来讲讲REGISTER_OP宏背后的具体原理。我们在上面已经给出了此宏的定义,此处针对它的实现展开谈谈: // 先来看op.h #define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name) #define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name) #define REGISTER_OP_UNIQ(ctr, name) static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr TF_ATTRIBUTE_UNUSED = ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( name)>(name) // 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(); }); // op定义的链式规则是通过OpDefBuilderWrapper类实现的 class OpDefBuilderWrapper<true> { public: explicit OpDefBuilderWrapper(const char name[]) : builder_(name) {} OpDefBuilderWrapper<true>& Input(string spec) { builder_.Input(std::move(spec)); return *this; // 显而易见,调用Input仍然返回OpDefBuilderWrapper<true>本身 } OpDefBuilderWrapper<true>& Output(string spec) { builder_.Output(std::move(spec)); return *this; } OpDefBuilderWrapper<true>& SetShapeFn( Status (*fn)(shape_inference::InferenceContext*)) { builder_.SetShapeFn(fn); return *this; } const ::tensorflow::OpDefBuilder& builder() const { return builder_; } private: mutable ::tensorflow::OpDefBuilder builder_; }; // 当通过链式规划构建好op后,再通过OpDefBuilderReceiver完成op的注册 // op.h struct OpDefBuilderReceiver { // To call OpRegistry::Global()->Register(...), used by the // REGISTER_OP macro below. // Note: These are implicitly converting constructors. OpDefBuilderReceiver( const OpDefBuilderWrapper<true>& wrapper); // NOLINT(runtime/explicit) constexpr OpDefBuilderReceiver(const OpDefBuilderWrapper<false>&) { } // NOLINT(runtime/explicit) }; // op.cc,然后在OpDefBuilderReceiver构造函数内部完成OpDefBuilderWrapper的全局注册 OpDefBuilderReceiver::OpDefBuilderReceiver( const OpDefBuilderWrapper<true>& wrapper) { OpRegistry::Global()->Register( [wrapper](OpRegistrationData* op_reg_data) -> Status { return wrapper.builder().Finalize(op_reg_data); }); }
四、总结
本文为大家系统讲解了TensorFlow的核心抽象op及其kernel实现。需要自定义op的具体场景,以及op的运行框架及若干技术细节。读罢此文,读者应该有如下几点收获:
参考资料
1.《TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems》: https://arxiv.org/abs/1603.04467
2.Graphs and Sessions: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/graphs.md
3.Adding a New Op: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/extend/op.md
4.跨设备通信send/recv: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/sendrecv_ops.h
5.OpKernel definition: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h
6.tensorflow源码解析之framework-resource: https://www.cnblogs.com/jicanghai/p/9535504.html
7.tensorflow源码解析之framework-op: https://www.cnblogs.com/jicanghai/p/9539513.html
**本文作者:李杰**
TF计算图从逻辑层来讲,由op与tensor构成。op是项点代表计算单元,tensor是边代表op之间流动的数据内容,两者配合以数据流图的形式来表达计算图。那么op对应的物理层实现是什么?TF中有哪些op,以及各自的适用场景是什么?op到底是如何运行的?接下来让我们一起探索和回答这些问题。
## 一、初识op
### 1.1 op定义
op代表计算图中的节点,是tf.Operation对象,代表一个计算单元。用户在创建模型和训练代码时,会创建一系列op及其依赖关系,并将这些op和依赖添加到tf.Graph对象中(一般为默认图)。比如:tf.matmul()就是一个op,它有两个输入tensor和一个输出tensor。
### 1.2 op分类
op的分类一般有多个视角,比如按是否内置划分、按工作类型划分。
按是否内置划分,一般分为:内置op和自定义op(见“二、自定义op”部分介绍)。
按工作类型划分,一般分为:常见数学op、数组op、矩阵op、有状态op、神经网络op、检查点op、队列与同步op、控制流op。TF白皮书对内置op的分类总结如下:
![](https://p3-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/2ba4e499f2ec4c869067ccf14cfc41c7~tplv-k3u1fbpfcp-zoom-1.image)
### 1.3 op与kernel
op一般都有名称且代表一个抽象的计算过程。op可以设置若干属性,但这些属性必须在编译期提供或推理得到,因为它们用来实例化一个节点对象从而执行真正的计算。属性的经典用法就是拿来支持类型多态,比如两个浮点张量的矩阵乘法与两个整型张量的矩阵乘法。
kernel是op在指定设备类型(CPU/GPU)上的具体实现。TF二进制库通过注册机制定义了一系列op及对应的kernel实现,用户可以提供额外的op定义与kernel实现进行扩充。一般来说,一个op对应多个kernel实现。
接下来让我们一起用矩阵乘法MatMul算子的相关代码来理解op与kernel的关系(此处不必纠结代码细节,只需体会op与kernel关系即可):
```// 首先给出op注册的定义。其中输入输出支持泛型,其合法类型在Attr中进行枚举。// 代码位置 tensorflow1.15.5tensorflowcoreopsmath_ops.ccREGISTER_OP("MatMul") .Input("a: T") .Input("b: T") .Output("product: T") .Attr("transpose_a: bool = false") .Attr("transpose_b: bool = false") .Attr( "T: {bfloat16, half, float, double, int32, int64, complex64, " "complex128}") .SetShapeFn(shape_inference::MatMulShape); // MatMul的实现,采用类模板机制// 代码位置 tensorflow1.15.5tensorflowcorekernelsmatmul_op.cctemplate <typename Device, typename T, bool USE_CUBLAS>class MatMulOp : public OpKernel { public: explicit MatMulOp(OpKernelConstruction* ctx) : OpKernel(ctx), algorithms_set_already_(false) { OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
LaunchMatMul<Device, T, USE_CUBLAS>::GetBlasGemmAlgorithm( ctx, &algorithms_, &algorithms_set_already_); use_autotune_ = MatmulAutotuneEnable(); } // 省略了很多代码... private: std::vector<int64> algorithms_; bool algorithms_set_already_; bool use_autotune_; bool transpose_a_; bool transpose_b_;};
// MatMul的op定义与kernel实现绑定处理// 代码位置 tensorflow1.15.5tensorflowcorekernelsmatmul_op.cc#define REGISTER_CPU_EIGEN(T) /*cpu与eigen组合对应实现*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>);
#define REGISTER_CPU(T) /*cpu对应实现(eigen与非eigen)*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T"), MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); REGISTER_CPU_EIGEN(T);
#define REGISTER_GPU(T) /*gpu对应实现(cublas与非cublas)*/ REGISTER_KERNEL_BUILDER( Name("MatMul").Device(DEVICE_GPU).TypeConstraint<T>("T"), MatMulOp<GPUDevice, T, true /* cublas, true by default */>); REGISTER_KERNEL_BUILDER(Name("MatMul") .Device(DEVICE_GPU) .TypeConstraint<T>("T") .Label("cublas"), MatMulOp<GPUDevice, T, true /* cublas */>)```
## 二、自定义op
用户编写的模型训练代码一般由TF原生的op算子及其依赖关系组成,但有时候我们定义的计算逻辑在TF中没有相应的op实现。根据TensorFlow官网的建议,我们应当先组合python op算子或python函数进行尝试。完成尝试之后再决定要不要自定义op。
### 2.1 自定义op场景
一般来说,需要自定义op的场景有如下3个:
•用TF原生op组合来表达新计算逻辑的过程比较复杂或不可能
•用TF原生op组合来表达新计算逻辑,其计算性能较低
•在新版编译器中也较难实现op融合的计算逻辑需要我们手动实现融合
在此举个例子方便大家理解。假如我们要实现一个新计算实逻:中位数池化(median pooling),过程中要在滑动窗口不断求得中位数。检索TF文档没有发现对应op,因此我们先考虑用TF python op组合来实现它,果然通过**ExtractImagePatches** and **TopK**就可以实现这个功能。经测试前述组合方案并不是计算和存储高效的,因此我们就有必要将median pooling在一个op中进行高效实现。
### 2.2 自定义op流程
自定义op一般遵循5个基本步骤:
1.注册op,具体包括:指定名称、输入/输出声明、形状函数。
2.定义kernel(即op的实现)并与op绑定。一个op有多个kernel实现,具体由输入输出类型、硬件(CPU、GPU)决定。
3.创建python包装器,一般由op注册机制自动完成。
4.编写op的梯度计算函数(可选项)。
5.测试op,通过python测试较为方便,当然也可通过C++进行测试。
接下来我们就以官网最简单的ZeroOut同步式自定义op(继承OpKernel)为例,结合代码来讲述上述5个步骤。下面先给出步骤1和步骤2用C++实现的代码(官方推荐用bazel编译so文件):
```// 步骤1:注册opREGISTER_OP("ZeroOut").Input("to_zero: int32").Output("zeroed: int32").SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); //c's input and output type is std::vector<ShapeHandle> return Status::OK(); });
// 步骤2:定义kernel(常规CPU设备),并把kernel与op绑定class ZeroOutOp : public OpKernel {public: explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override { // Grab the input tensor from OpKernelContext instance 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)); // OP_REQUIRES_OK第二个参数一般为方法调用,此处为输出张量分配内存空间 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); }};
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);```
步骤3加载上述so文件(自动完成前后端op映射);步骤4是可选项,此处不需要;步骤5基于python api测试op功能。相应代码如下:
```import tensorflow as tfzero_out_module = tf.load_op_library('./zero_out.so') # 加载so文件生成python modulewith tf.Session(''): zero_out_module.zero_out([[1, 2], [3, 4]]).eval()
# Printsarray([[1, 0], [0, 0]], dtype=int32)```
### 2.3 高级话题
关于op的技术话题还有很多,我们在此简述一些要点:
1.如果实现了一个多线程CPU kernel,则可以利用work_sharder.h中的Shard函数。
2.大多数op以同步方式工作,只需继承OpKernel改写Compute()方法,且此方法必须线程安全。
3.如果一个op因为其它op的运行而阻塞,则这个op可以采用异步方式工作,继承AsyncOpKernel改写ComputeAsync()方法,且此方法必须线程安全。异步op最经典的例子就是跨设备通信send/recv pair中的RecvOp。
4.如果要为op配置一些静态属性,可使用Attr,它有一套特有的支持类型。典型应用是支持泛型。
5.实现GPU kernel有两部分内容:OpKernel和CUDA kernel,相应的加载代码。
6.编译自定义op,首先要配置头文件搜索路径与库文件搜索路径,接着指定编译和链接选项,最后还要确保ABI兼容性。
7.Resource(资源)代表相同设备上op共享的内容,比如:张量值、kv存储表、队列、读取器、网络连接等。代表资源的类必须继承ResourceBase,然后注册ResourceHandleOp生成资源句柄,普通op以resouce类型的Input进行引入。
## 三、op工作原理
### 3.1 op运行框架
整体来看,op与kernel都有其结构描述与统一的注册管理中心。而OpDefBuilder有两个包装类OpDefBuilderWrapper和OpDefBuilderReceiver,前者支持op构建的链式语法,后者接受op构建结果并进行注册。众所周知,op是编译期概念,而kernel是运行期概念,在AI编译器的后端处理流程中会进行op的算子选择,此过程会基于一系列策略为op匹配最合适的kernel实现。
![](https://p3-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/419cd70db0da46b090f740d7245ee836~tplv-k3u1fbpfcp-zoom-1.image)
### 3.2 若干技术细节
首先,我们来看一下大家在使用TensorFlow过程中经常碰到的libtensorflow_framework.so。按照tf1.15.5/tensorflow/BUILD中的描述,libtensorflow_framework.so定义了op和kernel的注册机制而不涉及具体实现。
```// rootdir=tensorflow1.15.5// ${rootdir}/tensorflow/BUILD/*# A shared object which includes registration mechanisms for ops and# kernels. Does not include the implementations of any ops or kernels. Instead,# the library which loads libtensorflow_framework.so# (e.g. _pywrap_tensorflow_internal.so for Python, libtensorflow.so for the C# API) is responsible for registering ops with libtensorflow_framework.so. In# addition to this core set of ops, user libraries which are loaded (via# TF_LoadLibrary/tf.load_op_library) register their ops and kernels with this# shared object directly.*/tf_cc_shared_object( name = "tensorflow_framework", framework_so = [], linkopts = select({ "//tensorflow:macos": [], "//tensorflow:windows": [], "//tensorflow:freebsd": [ "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)", "-lexecinfo", ], "//conditions:default": [ "-Wl,--version-script,$(location //tensorflow:tf_framework_version_script.lds)", ], }), linkstatic = 1, per_os_targets = True, soversion = VERSION, visibility = ["//visibility:public"], deps = [ "//tensorflow/cc/saved_model:loader_lite_impl", "//tensorflow/core:core_cpu_impl", "//tensorflow/core:framework_internal_impl", /* 展开此target进行查看 */ "//tensorflow/core:gpu_runtime_impl", "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry_impl", "//tensorflow/core:lib_internal_impl", "//tensorflow/stream_executor:stream_executor_impl", "//tensorflow:tf_framework_version_script.lds", ] + tf_additional_binary_deps(),)
// ${rootdir}/tensorflow/core/BUILDtf_cuda_library( name = "framework_internal_impl", srcs = FRAMEWORK_INTERNAL_PRIVATE_HEADERS + glob( // 可以查看FRAMEWORK_INTERNAL_PRIVATE_HEADERS内容 [ "example/**/*.cc", "framework/**/*.cc", "util/**/*.cc", "graph/edgeset.cc", "graph/graph.cc", "graph/graph_def_builder.cc", "graph/node_builder.cc", "graph/tensor_id.cc", "graph/while_context.h", "graph/while_context.cc", ], // 省略了诸多代码)
// FRAMEWORK_INTERNAL_PRIVATE_HEADERS的内容FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [ "graph/edgeset.h", "graph/graph.h", "graph/graph_def_builder.h", "graph/node_builder.h", "graph/tensor_id.h",] + glob( [ "example/**/*.h", "framework/**/*.h", // 这里就是重点,查看${rootdir}/tensorflow/core/framework/op.h和opkernel.h "util/**/*.h", ])
// 先来看op.h#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)#define REGISTER_OP_UNIQ(ctr, name) static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr TF_ATTRIBUTE_UNUSED = ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( name)>(name) // 再来看看opkernel.h#define REGISTER_KERNEL_BUILDER(kernel_builder, ...) REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__)
#define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)
#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...) constexpr bool should_register_##ctr##__flag = SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__); static ::tensorflow::kernel_factory::OpKernelRegistrar registrar__body__##ctr##__object( should_register_##ctr##__flag ? ::tensorflow::register_kernel::kernel_builder.Build() : nullptr, #__VA_ARGS__, [](::tensorflow::OpKernelConstruction* context) -> ::tensorflow::OpKernel* { return new __VA_ARGS__(context); });```
参照上述同样的流程,我们可以发现libtensorflow.so中涉及op与kernel的具体实现,同时也包括Session的具体实现。
最后,我们再来讲讲REGISTER_OP宏背后的具体原理。我们在上面已经给出了此宏的定义,此处针对它的实现展开谈谈:
```// 先来看op.h#define REGISTER_OP(name) REGISTER_OP_UNIQ_HELPER(__COUNTER__, name)#define REGISTER_OP_UNIQ_HELPER(ctr, name) REGISTER_OP_UNIQ(ctr, name)#define REGISTER_OP_UNIQ(ctr, name) static ::tensorflow::register_op::OpDefBuilderReceiver register_op##ctr TF_ATTRIBUTE_UNUSED = ::tensorflow::register_op::OpDefBuilderWrapper<SHOULD_REGISTER_OP( name)>(name)
// 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(); });
// op定义的链式规则是通过OpDefBuilderWrapper类实现的class OpDefBuilderWrapper<true> { public: explicit OpDefBuilderWrapper(const char name[]) : builder_(name) {}
OpDefBuilderWrapper<true>& Input(string spec) { builder_.Input(std::move(spec)); return *this; // 显而易见,调用Input仍然返回OpDefBuilderWrapper<true>本身 } OpDefBuilderWrapper<true>& Output(string spec) { builder_.Output(std::move(spec)); return *this; }
OpDefBuilderWrapper<true>& SetShapeFn( Status (*fn)(shape_inference::InferenceContext*)) { builder_.SetShapeFn(fn); return *this; } const ::tensorflow::OpDefBuilder& builder() const { return builder_; }
private: mutable ::tensorflow::OpDefBuilder builder_;};
// 当通过链式规划构建好op后,再通过OpDefBuilderReceiver完成op的注册// op.hstruct OpDefBuilderReceiver { // To call OpRegistry::Global()->Register(...), used by the // REGISTER_OP macro below. // Note: These are implicitly converting constructors. OpDefBuilderReceiver( const OpDefBuilderWrapper<true>& wrapper); // NOLINT(runtime/explicit) constexpr OpDefBuilderReceiver(const OpDefBuilderWrapper<false>&) { } // NOLINT(runtime/explicit)};
// op.cc,然后在OpDefBuilderReceiver构造函数内部完成OpDefBuilderWrapper的全局注册OpDefBuilderReceiver::OpDefBuilderReceiver( const OpDefBuilderWrapper<true>& wrapper) { OpRegistry::Global()->Register( [wrapper](OpRegistrationData* op_reg_data) -> Status { return wrapper.builder().Finalize(op_reg_data); });}```
## 四、总结
本文为大家系统讲解了TensorFlow的核心抽象op及其kernel实现。需要自定义op的具体场景,以及op的运行框架及若干技术细节。读罢此文,读者应该有如下几点收获:
•TensorFlow中op是编译期概念,kernel是运行期概念,两者各自的定义与注册方式,以及相应的映射逻辑。
•掌握TensorFlow的高阶玩法:自定义op。这将使你之前工作的不可能变为可能,由低效转化为高效。
•掌握op与kernel注册的宏定义来自何方,以及宏定义背后具体的运行框架。
## 参考资料
1.《TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems》: <https://arxiv.org/abs/1603.04467>
2.Graphs and Sessions: <https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/graphs.md>
3.Adding a New Op: <https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/extend/op.md>
4.跨设备通信send/recv: <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/sendrecv_ops.h>
5.OpKernel definition: <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h>
6.tensorflow源码解析之framework-resource: <https://www.cnblogs.com/jicanghai/p/9535504.html>
7.tensorflow源码解析之framework-op: <https://www.cnblogs.com/jicanghai/p/9539513.html>
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:Dive into TensorFlow系列(2)- 解析TF核心抽象op算子 - Python技术站