使用Tensorflow C++ API自定义操作

2024年1月3日 234次阅读 来源: Kapok

Tensorflow提供了大量的基本操作使得我们能够任意组合来实现我们需要的复杂操作，但有时候我们需要的操作不太容易通过这些基本操作来组合，或者复杂的组合方式带来严重的性能开销，这时我们可能会考虑去使用py_func来包装Python函数借助Numpy来实现，但性能方面可能也无法达到满意的程度，更不要说有些操作不适合向量化的写法，这个时候用C++ API来实现自己的一个操作可能会是更好的选择。

总得来说，我们可以通过静态链接或动态连接的方式来添加我们自定义的操作，前者需要你能成功从源码编译安装Tensorflow（Installing TensorFlow from Sources），后者只需要安装Tensorflow的Python包即可，不过后者有一个缺点就是可用的C++ API是有限的，如果你需要一些别的接口可能要自己去找相关的头文件然后手动添加进来，如果头文件数目很多这就是一件很头疼的事情了，但优点也是很明显，一是简单，二是不用强制要求使用你操作的人也要从源码编译安装Tensorflow。下面主要以动态链接的添加为例进行叙述，对静态链接添加的方式也会有涉及，二者基本上是一样的。另外，本文主要注重实践部分，如果你需要一些更详细的标准化说明，可以去参考官方的Tutorial：Adding a New Op，一些比如Ops注册的属性说明以及如何保持Ops属性的后向兼容性等会更加详细。

另外如果你对Tensorflow注册Ops和Kernel的过程或者它们是如何被Python代码调用的感兴趣，可以阅读我之前写的一篇文章Tensorflow是如何注册和调用C++ New Op的。如果你想学习一下官方Ops是怎么写的，我推荐从Bias_Op来开始学习，因为这个Op写的还是很清晰然后对初学者不会有太多的阻碍，当然也可以参考我最近写的PSROIAlign。这里还必须要推荐一个利器SourceInsight简直是阅读这种开源代码的法宝。

先总结官方给的自定义op的标准流程：注册Op，实现Op，创建python接口，实现Op梯度计算（如果不需要求导也可以直接pass掉，实现可以在python端也可以用py_func去包装其他python函数，也可以再写一个C++ Op来专门计算梯度），测试。

  • 注册Op

注册Op相当于是一个声明的过程。Op是tensorflow中非常重要的概念，一个Op接收一个或多个输入张量，然后经过某种运算，产生其他零个或多个tensor，然后这些tensor又可以被其他Op使用。类似于C++中我们定义变量需要知道数据类型，字节数等信息一样，创建一个Op同样需要一些额外信息包括attributes（输入输出类型以及合法取值等，也可以看作是Op的输入但是不同于输入的是属性永远是常量，其值在Op被添加到图中时被设置，并且是一直放在CPU上的）以及输入输出列表，还可以直接加Doc，具体信息是我们在REGISTER_OP时指定的，REGISTER_OP是一个宏，其内部实现是一个wrapper利用了C++中的常用伎俩chaining调用实现，所有你在这添加的信息都会以另一种形式出现在动态生成的Python代码中。有一点需要注意，在C++这边Ops的名称必须是CamelCase类型的，在Python那一边会自动被转换成Python风格的snake_case类型。

注册这个地方还有一个SetShapeFn需要说一下，主要作用是检查输入的shape并指定输出的shape，当然你也可以在Op的compute里面检查之类，但是这个ShapeFn有一个点是可以让tesorflow不用执行操作就能获取输入输出信息。在ShapeFn里面你可以拿到输入输出的每一个维度的大小（DimensionHandle），或者属性常量的值或者输入常量的值，然后组合成输出的ShapeHandle，最后调用set_output指定对应输出的shape，同时DimensionHandle是可以做四则运算的。一开始我对如何指定输出大小的api也有一些困惑，因为还涉及动态shape，这里推荐仔细阅读InferenceContext这个类，还是推荐用SourceInsight用好搜索和bookmark即可。

  • 实现Op

动手之前了解一下C++中的functor、模板及其特化还是很有必要的，对lambda也有了解的话就更好了，如果你对C++不熟的话建议尽量避免使用Eigen，直接把数组取出来用C计算就行，因为tensorflow里面的张量都是按行主序存的（多维的情况就是最外面的那一维变化最快）。

用C++实现Op有一个固定的套路，遵循这个套路可以避免走弯路，当然这都不是必须的，只要你定义了计算函数并且在kernel的Compute里面调用你的计算即可。初学者可以参考下面这个框架来做：

  1. 定义一个Functor模板类做实际的计算工作
  2. 针对不同设备甚至不同数据类型特化Functor模板类
  3. 定义一个Kernel模板类，继承自OpKernel，在构造函数中根据传进来的OpKernelConstruction设置必要的成员
  4. 视情况针对别的设备特化Kernel模板类，一般无需特化，因为这个类里面一般会做一些通用工作，然后将实际的计算转到Functor模板类中
  5. 重写Kernel的Compute方法，利用OpKernelContext获取输出，分配输出，并进行合法性检查，然后转调对应的计算Functor
  6. 注册Kernel

很多重要的api都在OpKernelConstruction和OpKernelContext两个类里面建议详细阅读。同时Compute方法必须是线程安全的，因此任何对类成员的访问必须要用互斥保护。

在转到实际的计算函数前通常会把输入输出Tensor的缓冲区取出来，要么变成Eigen的表示即TensorMap（其对应的很多成员要去TensorBase里面去找），要么更进一步直接再调用TensorMap的data方法把缓冲区指针取出来传给计算函数。具体地，可以去看Tensor这个类提供的一些接口。

计算的实现过程不细说，可以直接上C，可以用std::thread保证移植性，也可以利用强大的Eigen，里面有很好的并行化机制根据每个执行单元的cost来分配线程资源，还可以利用Tensorflow提供的对Eigen进行包装后的Shard（这个头文件需要自己加）工具类，下面的代码是一个使用示例：

auto work_routine = [&your_capture](int64_t start, int64_t limit){
 for (int64_t worker_index = start; worker_index < limit; ++worker_index){
 // do something  }
};

const DeviceBase::CpuWorkerThreads& worker_threads = *(context->device()->tensorflow_cpu_worker_threads());
const int64_t shard_cost = 4 * num_rois;
Shard(worker_threads.num_threads, worker_threads.workers, total_elems, shard_cost, work_routine);

GPU的实现多数简单情况下GetCudaLaunchConfig + CUDA_1D_KERNEL_LOOP就可以搞定，更多有用的接口可以去看cuda_kernel_helper.h，也可以借助cub的api或者更暴力一点直接写。顺便说一下，使用Eigen或者Tensorflow的CPU并行化机制编写代码，通常写出的代码整体逻辑和GPU代码基本一致，我在写PSROIAlign的时候从CPU代码移植到GPU改动的地方很少。

OpKernel的注册和Ops的注册比较类似，也是调用一个宏REGISTER_KERNEL_BUILDER，指定名称、Kernel对应的设备类型，以及创建这个Kernel的C++类。详见Tensorflow是如何注册和调用C++ New Op的。

  • 创建python接口

在这之前当然先要对自己写的Op进行编译，我捣鼓了一个CMakeLists感觉很好用，可以把C++和CUDA代码分开编译然后一起链接很省事，只要正确安装了cuda和python包头文件都可以自己找到，推荐给大家CMakeLists.txt。

编译成功后应该可以获得一个动态库文件 .so，python这边load一下然后包装一下就好了，把官方的示例抄过来：

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out
  • 实现Op梯度计算

一般就是python这边要么用tensorflow自带Op进行组合，要么再去调用另一个计算梯度的自定义Op，然后整个计算过程放在一个函数里面，用@ops.RegisterGradient修饰一下就行了，具体可参见官方文档。

  • 测试Op

同样官方有示例，推荐一个tf.test里面的compute_gradient和compute_gradient_error，很好用。注意这两个函数计算的并不是梯度，而是输出对输入的Jacobian，计算梯度值要用tf.gradients，tf.test中的那两个函数也是基于tf.gradients。其中compute_gradient来计算理论和数值Jacobian，compute_gradient_error计算二者之间的误差。

  • 总结

最后总结一下我目前遇到的坑：

  1. 最好不要用Eigen接口，除非你对Eigen比较熟，否则建议使用前阅读以下Lazy Evaluation and Aliasing和Eigen: Common pitfalls
  2. 用Shard工具类写CPU端的kernel，方便移植到GPU上，但要注意线程间的同步
  3. 全局一致的常量输入使用op的Attr来指定，尤其是需要基于这些常量输入做进一步地运算的时候，因为如果将常量输入作为Scalar类型的Tensor输入，那么在CPU上和GPU上运行时这些输入将在不同的内存里，如果要基于这些常量做进一步地运算在GPU上要用cuda kernel，不利于代码结构的简化
  4. 输出记得先清零

此外，如果你希望静态链接你的Op，那么可以把代码放在Tensorflow源代码的user_ops目录下，然后在BUILD文件里添加一个tf_custom_op_library就可以了。

暂时想到的就这么多，最后祝大家炼丹顺利~~找到理想工作。

    原文作者：Kapok
    原文地址: https://zhuanlan.zhihu.com/p/34169502
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞