如何在张量流中编写自定义池模块？

Question

我认为添加一个新的模块将进行中心池化.

我正在查看tensorflow代码,并且有一个名为gen_nn_ops.py内部调用来自另一个文件的函数的文件,通过传递"Maxpool","AvgPool"等参数来执行所需的计算.

我想做中心池,选择窗口中的中心元素.我已为matlab和c ++版本准备好代码,但需要知道如何在TensorFlow中添加新模块以进行此类计算.还在哪里设置backpropogation代码.

Answer 1

自定义池层可能会在C++中实现.要了解您需要做什么,让我们看看tf.nn.max_pool_with_argmax()生活的实施地点:

1. Python包装器函数(tf.nn.max_pool_with_argmax()本身)是自动生成的gen_nn_ops.py.这最终导入nn.py,以便它出现在tf.nn你的时候import tensorflow as tf.

2. 在C++中,有一个op注册ops/nn_ops.cc,一个内核注册kernels/maxpooling_op.cc.

3. 该内核本身的定义kernels/maxpooling_op.cc.

4. 渐变被定义为"MaxPoolWithArgmaxGrad"在相同位置的单独操作.

添加新操作需要做大量的工作(有关更完整的指南,请参阅本教程),但希望这些指针可以提供帮助!