江苏龙网

  1. 首页 > 生活 > >

Tensorflow 中咋定义自己的层呢

最近把回答整理成了一个小短文：https://zhuanlan.zhihu.com/p/34169502
------------------------------
碰巧看到这个问题，前段时间做比赛顺便学习了一下用C++自定义tensorflow的op，比赛虽然做的不怎么样，但是学了不少额外的东西。。。
今天顺便总结一下，希望对初次接触的同学有所帮助。这里仅涉及动态链接的op添加，因为当时我使用的bazel有问题我也懒得降版本，导致无法从源码进行编译。动态添加是最简单的方式，只要安装了tensorflow的python包就可以编译然后链接使用，有一个缺点就是可用的api是有限的。
讲真的其实用C++写一个tensorflow的新操作还是很简单的，甚至只要会C然后利用一个好的框架就大致可以搞定，总的来说可以参考的资料主要有：tensorflow offical add new op 以及tensorflow源码中官方的op是怎么添加的，其实都是一种套路，多学习官方代码里的op事半功倍，具体来说我推荐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的输入但是不同于输入的是属性永远是常量）以及输入输出列表，还可以直接加Doc，具体信息是我们在REGISTER_OP时指定的，REGISTER_OP是一个宏，其内部实现是一个wrapper利用了C++中的常用伎俩chaining调用实现。具体注册的过程可以参看我之前写的一个简单的介绍 -- 用户：tensorflow中的op注册的原理是什么？
注册这个地方还有一个SetShapeFn需要说一下，主要作用是检查输入的shape并指定输出的shape，当然你也可以在op的compute里面检查之类，但是这个ShapeFn有一个点是可以让tesorflow不用执行操作就能获取输入输出信息。一开始我对如何指定输出大小的api也有一些困惑，因为还涉及动态shape，这里推荐仔细阅读InferenceContext这个类，还是推荐用SourceInsight用好搜索和bookmark即可。
实现op
动手之前了解一下C++中的functor、模板及其特化还是很有必要的，对lambda也有了解的话就更好了，如果你对C++不熟的话建议尽量避免使用Eigen，直接把数组取出来用C计算就行，因为tensorflow里面的张量都是按行主序存的（多维的情况就是最外面的那一维变化最快）。
用C++实现op有一个固定的套路，遵循这个套路可以避免走弯路，当然这都不是必须的，只要你定义了计算函数并且在kernel的Compute里面调用你的计算即可。初学者可以参考下面这个框架来做：
Tensorflow 中咋定义自己的层呢

每个Op继承自OpKernel，然后最后在注册kernel的时候可以对不同的设备类型进行特化（如果你CPU和GPU实现相差很多，那么最好是定义一个对多数设备通用的类模板，然后对特殊设备特化这个模板）。OpKernel::Compute()是每个op计算的入口，每个op都要进行重写，通常取出所有的输入做一些常规检查，然后分配输出和临时tensor，最后转去调用实际的计算函数。同样很多重要的api都在OpKernelConstruction这个类里面建议详细阅读。


推荐阅读


上一篇：中国移动采购共享服务中心，有人了解这个中心的情况么，可以介绍一下吗

下一篇：现在学it还有用吗