Tensorflow, How can I compute backward pass for a given forward function - tensorflow

Question

I want to construct an L2-norm layer in Caffe style (well, I actually want to use Tensorflow in a pycaffe layer, since using CUDA to write .cu files in Caffe is an onerous task.)
Forward pass:
- input(x): n-D array
- output(y): n-D array that has the same shape of input
- operation:
y = x / sqrt(sum(x^2,axis=(0,1))) # channel wise L2 normalization
class L2NormLayer:
def __init__(self):
self.eps = 1e-12
self.sess = tf.Session()
def forward(self, in_x):
self.x = tf.constant(in_x)
self.xp2 = tf.pow(self.x, 2)
self.sum_xp2 = tf.reduce_sum(self.xp2, axis=(0, 1))
self.sqrt_sum_xp2 = tf.sqrt(self.sum_xp2 + self.eps)
self.hat = tf.div(self.x, self.sqrt_sum_xp2)
return self.sess.run(self.hat)
def backward(self, dl):
# 'dl' is loss calculated at upper layer (chain rule)
# how do I calculate this gradient automatically using Tensorflow
# hand-craft backward version
loss = tf.constant(dl)
d_x1 = tf.div(loss, self.sqrt_sum_xp2)
d_sqrt_sum_xp2 = tf.div(-tf.reduce_sum(self.x * dl, axis=(0, 1)), (self.eps + tf.pow(self.sqrt_sum_xp2, 2)))
d_sum_xp2 = tf.div(d_sqrt_sum_xp2, (self.eps + 2 * tf.sqrt(self.sum_xp2)))
d_xp2 = tf.ones_like(self.xp2) * d_sum_xp2
d_x2 = 2 * self.x * d_xp2
d_x = d_x1 + d_x2
return self.sess.run(d_x)
As commented in the code, how can I calcualte the gradient of the forward pass function by using Tensorflow automatically?

Answer 1

I think your best strategy would be to use existing caffe layers to achieve your goal.
First, use "Reduction" layer to compute the sq. L2 norm of x:
layer {
name: "norm_x_sq"
type: "Reduction"
bottom: "x"
top: "norm_x_sq"
reduction_param { operation: SUMSQ axis: 1 }
}
Use "Power" layer to take the square root of the norm and compute its reciprocal:
layer {
name: "norm_x-1"
type: "Power"
bottom: "norm_x_sq"
top: "norm_x-1"
power_param { power: -0.5 }
}
Once you have the denominator, you need to "Tile" it back to the same shape as x:
layer {
name: "denom"
type: "Tile"
bottom: "norm_x-1"
top: "denom"
tile_param { axis:1 tiles: N } # here you'll have to manually put the target dimension N
}
Finally, use "Eltwise" layer to normalize x:
layer {
name: "x_norm"
type: "Eltwise"
bottom: "x"
bottom: "denom"
top: "x_norm"
eltwise_param { operation: PROD }
}
Some additional notes:
1. Dividing by the norm might be numerically unstable if the norm is very little. You might want to consider adding a tiny constant to "norm_x_sq" before taking the reciprocal of the square root. You can do that using existing layers as well.
2. This example showed how to normalize according to axis=1 dimension. Depending how your vectors are arranged in the blob, you might be able to use "Scale" layer for the division instead of tile+eltwise.
3. You might also find this thread useful.