def loop_fn(i):
with g:
x1 = array_ops.gather(x, i)
output = nn.max_pool3d(
x1, ksize, strides=strides, padding="VALID", data_format="NDHWC")
loss = nn.l2_loss(output)
ones = array_ops.ones_like(output)
g.watch(ones)
grad = g.gradient(loss, x1, output_gradients=ones)
grad_grad = g.gradient(grad, ones)
return output, grad, grad_grad
def loop_fn(i):
with g:
x1 = array_ops.gather(x, i)
output = nn.max_pool3d(
x1, ksize, strides=strides, padding="VALID", data_format="NDHWC")
loss = nn.l2_loss(output)
ones = array_ops.ones_like(output)
g.watch(ones)
grad = g.gradient(loss, x1, output_gradients=ones)
grad_grad = g.gradient(grad, ones)
return output, grad, grad_grad
def testMaxPool3DValidPaddingWithStridesF32(self):
with self.session() as sess:
with ops.device("/device:IPU:0"):
pa = array_ops.placeholder(np.float32, [1, 1, 10, 10, 10], name="a")
output = nn.max_pool3d(pa,
ksize=[1, 1, 5, 5, 5],
strides=[1, 1, 2, 2, 2],
data_format='NCDHW',
padding='VALID',
name="max")
fd = {pa: np.ones([1, 1, 10, 10, 10])}
result = sess.run(output, fd)
self.assertAllClose(result, np.ones([1, 1, 3, 3, 3]))