How to write a custom pooling layer module in tensor flow?

Question

I have thought that adding a new module which will do the center pooling.

I was looking in the tensorflow code and there is a file named gen_nn_ops.py which internally calls function from another file by passing "Maxpool", "AvgPool", etc. parameters to do the required computation.

I want to do centre pooling which selects the center element in the window. I have code ready for the matlab and c++ versions but need to know how to add a new module in TensorFlow for such computation. Also where to set the backpropogation code.

Answer 1

A custom pooling layer would probably be implemented in C++. To see what you'd need to do, let's see where the implementation of tf.nn.max_pool_with_argmax() lives:

1. The Python wrapper function (tf.nn.max_pool_with_argmax() itself) is automatically generated, in gen_nn_ops.py. This is ultimately imported into nn.py, so that it appears under tf.nn when you import tensorflow as tf.

2. In C++, there is an op registration in ops/nn_ops.cc, and a kernel registration in kernels/maxpooling_op.cc.

3. The kernel itself is defined in kernels/maxpooling_op.cc.

4. The gradient is defined as a separate op—"MaxPoolWithArgmaxGrad"—in the same places.

There's a fair amount of work to do to add a new op (see this tutorial for a more complete guide), but hopefully these pointers can help!