Tensorflow上手3: 实现自己的Op

Tensorflow作为一个深度学习框架，主要强调将计算过程表示为数据流．计算图．它提供了大量的基本操作让我们可以任意组合，实现比普通的神经网络更强大的计算方法．但现实工作当中，我们常常会需要一些并不太容易实现的基本操作，比如我们之前说到的特殊的metric．有时候我们可以通过前面介绍的py_func来包装 Python 函数，但是在设计到性能和分布式训练的时候，我们就需要采用C++来实现自己的操作了．

于是今天我们就来谈谈自定义Op当中可能遇到的问题，按照Tensorflow的官方介绍分为：注册Op，实现Op，Python端接口和测试几个部分．

注册一个Op

声明一个Op的第一步就是采用宏 REGISTER_OP 注册Op．在采用 REGISTER_OP 注册的时候，通常会声明以下这些信息

REGISTER_OP(Op_name)
 .Attr("Name: type = default value")
 .Input("Name: type")
 .Output("Name: type")
 .SetShapeFn(a function)
 .Doc(R"doc(...)")

Attr 为Op的参数，既可以是单一的数值，也可以定义数组，比如 list(int)=[400, 600] ．

一个Op可以接收一个或者多个输入Tensor，然后产生零个或者多个输出Tensor，分别利用 Input 和 Output 定义．有两种情况比较特殊，一种是我们不确定输入和输出的类型，想支持多种类型时，可以定义template．

REGISTER_OP("ZeroOut")
 .Attr("T: {int, float}")
 .Input("to_zero: T")
 .Input("zeroed: T")

另一种是我们不确定输入和输出的个数，很抱歉，这个我也没有好的解决方法，通常我使用额外的attribute来表示特定的输入和输出是否有效．

另外需要特别强调的就是注册Op时的 SetShapFn 函数，他的主要作用是用来检查输入的shape是否满足指定形式，并且对输出的shape进行计算．这一点非常有用，能够让Tensorflow在定义计算图的时候就能获取输出信息，帮助判断给定的计算图是否合理．

在使用过程中，你可以通过 InferenceContext 中的 GetAttr , WithRank 来获得当前Op的attribute以及输入维度信息．可以通过 Divide , Multiply 将未知或已知的维度和常数进行计算，最后通过 set_output 对输出进行限制．在Tensorflow提供的image_ops.cc有很多很好的样例可以参考．

实现一个Op

在注册一个Op之后，就需要集成 OpKernel ，实现他的计算过程 Compute 函数，在 Compute 函数中，我们可以通过访问 OpKernelContext 来获得输入和输出信息．

先看Tensorflow的一个官方样例

#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的Op都采用Eigen来操作Tensor，可以通过 input_tensor.flat<int32>() 来获取Eigen向量，然后进行操作．

但我本人对C++并不非常熟练，也不了解Eigen，通常我会采用 input_tensor.flat<int32>().data() 来获取数组指针，进行操作．另外一方面我写的更多的是在GPU里面进行的计算，所以获得数组地址更有用一些．

在编写GPU上的Op的时候，有很多需要注意的地方，我个人体会最深的就是一定要避免使用 CudaMalloc 和 CudaFree 函数．这两个函数会Freeze GPU，降低训练和预测的效率．当我们需要申请新的内存地址时，可以通过 OpKernelContext 去申请 TempTensor 或者 PersistentTensor ．

Python接口与梯度实现

在Python当中使用新的Op很简单，只需要通过 tf.load_op_library 去调用C++生成的动态库文件就可以了．

这里稍微值得一提就是实现Op的梯度．因为Tensorflow构建的时计算图，所以梯度其实也是计算图的一部分，新Op的梯度可以是另一个新Op，也可以是现有Op组成的计算图，就比方Tensorflow官方提供的ZeroOut梯度．

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
 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]

函数的输入op和grad分别是当前op的相关信息以及输出Tensor的梯度．如果你的Op有两个或以上的输出，那么就需要定义多个梯度输入，如下

@ops.RegisterGradient("YourOp")
def _your_op(op, grad1, grad2):
 return something

有时候计算梯度需要输入向量，或者输出向量，那么我们可以通过op获得相关的，可以调用 op.inputs[i] , op.outputs[i] 和 op.get_attr ，比如

@ops.RegisterGradient("YourOp")
def _your_op(op, grad):
 weight = op.get_attr("weight")
 return weight * op.outputs[0] * \
 (1 - op.outputs[0]) * grad

在进行测试的时候，我们可以通过Tensorflow提供的 tf.test.compute_gradient_error 函数对梯度进行测试．他的工作原理是通过对输入的每一数进行微小的改变，计算 jacobian matrix ．这里有一个小坑，如果你实现的gradient和输出传进来的gradien有关，那么你有可能在测试时得到错误的结果，因为测试时传回来的gradient可能并不符合实际情况．比如对于一个本来应该忽略的数据传入梯度．

小结

从我个人的使用经验来说，实现Tensorflow的新Op是一个有意思但时而又枯燥的学习过程．他可以帮助你更好的理解Tensorflow的设计理念和实现细节，比如Tensorflow是如何保证ThreadSafe的，是如何进行内存管理的，一个Op的Compute和ComputeAsync又是如何实现的，计算图是如何将数据在GPU和CPU之间转换的等等．

这篇就到这里，之后我会准备学习和介绍更多关于Tensorflow计算图的概念和分布式训练的内容．