江苏龙网

  1. 首页 > 生活 > >

Tensorflow 中咋定义自己的层呢( 二 )


具体实现过程不细说了，可以直接上C，可以用std::thread保证移植性，也可以利用强大的Eigen，里面有很好的并行化机制根据各种cost来分配资源，还可以利用Tensorflow提供的Shard工具类，下面的代码是一个使用示例。GPU的实现我不在行，一般也就拿CUDA_1D_KERNEL_LOOP搞一下，tensorflow里面利用了大量cub的api很是优雅。顺便说一下，使用Eigen或者Tensorflow的CPU并行化机制编写代码，通常写出的代码整体逻辑和GPU代码基本一致，我在写PsRoiAlign的时候从CPU代码移植到GPU改动的地方很少。
const DeviceBase::CpuWorkerThreads\u0026amp; worker_threads = *(context-\u0026gt;device()-\u0026gt;tensorflow_cpu_worker_threads());const int64_t shard_cost = batch_size * num_rois * num_channals;Shard(worker_threads.num_threads, worker_threads.workers, pooled_features.size(), shard_cost, work_routine);创建python接口
在这之前当然先要对自己写的op进行编译，我捣鼓了一个CMakeLists感觉很好用，可以把C++和CUDA代码分开编译然后一起链接很省事，只要正确安装了cuda和python包头文件都可以自己找到，推荐给大家CMakeLists.txt。
编译成功后应该可以获得一个动态库文件 .so，python这边load一下然后包装一下就好了，把官方的示例抄过来：
import tensorflow as tfzero_out_module = tf.load_op_library(\u0026#39;./zero_out.so\u0026#39;)zero_out = zero_out_module.zero_out实现op梯度计算
一般我的做法就是python这边要么用tensorflow自带op进行组合，要么再去调用另一个计算梯度的自定义op，然后整个计算过程放在一个函数里面，用@ops.RegisterGradient修饰一下就行了，具体可参见官方文档。
测试：同样官方有示例，推荐一个tf.test里面的compute_gradient和compute_gradient_error，很好用。注意这两个函数计算的并不是梯度，而是输出对输入的Jacobian，计算梯度值要用tf.gradients，tf.test中的那两个函数也是基于tf.gradients。其中compute_gradient来计算理论和数值Jacobian，compute_gradient_error计算二者之间的误差。
最后总结一下我目前遇到的坑：
最好不要用Eigen接口，除非你对Eigen比较熟，否则建议使用前阅读以下Lazy Evaluation and Aliasing和Common pitfalls用Shard工具类写CPU端的kernel，方便移植到GPU上，但要注意线程间的同步全局一致的常量输入使用op的Attr来指定，尤其是需要基于这些常量输入做进一步地运算的时候，因为如果将常量输入作为Scalar类型的Tensor输入，那么在CPU上和GPU上运行时这些输入将在不同的内存里，如果要基于这些常量做进一步地运算在GPU上要用cuda kernel，不利于代码结构的简化输出记得先清零使用GetCudaLaunchConfig和CUDA_1D_KERNEL_LOOP暂时想到的就这么多，最后祝大家炼丹顺利~~找到理想工作。

■网友
在Eager Execution出来之前，可以使用py_func方法自定义层，但是可惜只能在CPU上运行。如果要在GPU上运行需要用C++编写。

在Eager Execution出来之后，可以轻松定义层和梯度，在定义时只要使用tf.matmul或者tf.multiply等tf的内置函数进行计算，自定义的层或者梯度就可以在GPU上执行。

可以参考Github上的TensorFlow Eager Execution示例教程：https://github.com/ZhuanZhiCode/TensorFlow-Eager-Execution-Examples

其中有一个自定义softmax操作，并为softmax自定义梯度的例子：https://github.com/ZhuanZhiCode/TensorFlow-Eager-Execution-Examples/blob/master/examples/CustomGradient/custom_softmax.py

注意，在Eager Execution机制中，你只要用tf内置操作定义前向传播，tf就会自动会你定义的函数计算梯度，但很多情况下，自动计算梯度的效率，不如你人工化简公示后自定义的梯度计算方法来的快。比较典型的例子就是Softmax，使用交叉熵的Softmax的梯度化简后就是object - target，如果用自动梯度，那就比这个复杂的多了。因此在很多情况下，自定义梯度是一个很重要的需求。


推荐阅读


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

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