Press "Enter" to skip to content

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,分别利用 InputOutput 定义.有两种情况比较特殊,一种是我们不确定输入和输出的类型,想支持多种类型时,可以定义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的时候,有很多需要注意的地方,我个人体会最深的就是一定要避免使用 CudaMallocCudaFree 函数.这两个函数会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之间转换的等等.

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注