def test_backward_grads_and_vars_channels_last(self):
"""Test `backward_grads` function with `channels_last` data format."""
with tf.device("/cpu:0"): # NHWC format
input_shape = (224, 224, 16)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
dy = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape,
data_format="channels_last")
dx_tr, grads_tr, vars_tr = residual.backward_grads_and_vars(
x, dy=dy, training=True)
dx_ev, grads_ev, vars_ev = residual.backward_grads_and_vars(
x, dy=dy, training=False)
self.assertNotAllClose(dx_tr, dx_ev)
self.assertTrue(isinstance(grads_tr, list))
self.assertTrue(isinstance(grads_ev, list))
self.assertTrue(isinstance(vars_tr, list))
self.assertTrue(isinstance(vars_ev, list))
for grad_tr, var_tr, grad_ev, var_ev in zip(
grads_tr, vars_tr, grads_ev, vars_ev):
if grad_tr is not None: # Batch norm moving mean, var gives None grad
self.assertEqual(grad_tr.shape, grad_ev.shape)
self.assertEqual(var_tr.shape, var_ev.shape)
self.assertEqual(grad_tr.shape, var_tr.shape)
def test_backward_grads_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (128, 8, 8)
data_shape = (16, ) + input_shape
# Use double precision for testing
x_true = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy = tf.random_normal(shape=data_shape, dtype=tf.float64)
residual = blocks._Residual(filters=128,
strides=(1, 1),
input_shape=input_shape,
fused=False,
dtype=tf.float64)
with tf.GradientTape() as tape:
x_true = tf.identity(x_true)
tape.watch(x_true)
y = residual(x_true, training=True)
# Gradients computed due to reversibility
x, dx, dw = residual.backward_grads(y, dy=dy, training=True)
vars_ = residual.trainable_variables
# True gradients computed by the tape
grads = tape.gradient(y, [x_true] + vars_, output_gradients=dy)
dx_true, dw_true = grads[0], grads[1:]
self.assertAllClose(x_true, x)
self.assertAllClose(dx_true, dx)
self.assertAllClose(dw_true, dw)
def test_backward_grads_and_vars_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (128, 8, 8)
data_shape = (16,) + input_shape
# Use double precision for testing
x_true = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy = tf.random_normal(shape=data_shape, dtype=tf.float64)
residual = blocks._Residual(
filters=128,
strides=(1, 1),
input_shape=input_shape,
fused=False,
dtype=tf.float64)
with tf.GradientTape() as tape:
x_true = tf.identity(x_true)
tape.watch(x_true)
y = residual(x_true, training=True)
# Gradients computed due to reversibility
x, dx, dw, vars_ = residual.backward_grads_and_vars(
y, dy=dy, training=True)
# True gradients computed by the tape
grads = tape.gradient(y, [x_true] + vars_, output_gradients=dy)
dx_true, dw_true = grads[0], grads[1:]
self.assertAllClose(x_true, x)
self.assertAllClose(dx_true, dx)
self.assertAllClose(dw_true, dw)
def test_backward_grads_and_vars_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
dy = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape)
dx_tr, grads_tr, vars_tr = residual.backward_grads_and_vars(
x, dy=dy, training=True)
dx_ev, grads_ev, vars_ev = residual.backward_grads_and_vars(
x, dy=dy, training=False)
self.assertNotAllClose(dx_tr, dx_ev)
self.assertTrue(isinstance(grads_tr, list))
self.assertTrue(isinstance(grads_ev, list))
self.assertTrue(isinstance(vars_tr, list))
self.assertTrue(isinstance(vars_ev, list))
for grad_tr, var_tr, grad_ev, var_ev in zip(
grads_tr, vars_tr, grads_ev, vars_ev):
if grad_tr is not None: # Batch norm moving mean, var gives None grad
self.assertEqual(grad_tr.shape, grad_ev.shape)
self.assertEqual(var_tr.shape, var_ev.shape)
self.assertEqual(grad_tr.shape, var_tr.shape)
def test_backward_grads_and_vars_channels_last(self):
"""Test `backward_grads` function with `channels_last` data format."""
with tf.device("/cpu:0"): # NHWC format
input_shape = (224, 224, 16)
data_shape = (16,) + input_shape
x = tf.random_normal(shape=data_shape)
dy = tf.random_normal(shape=data_shape)
residual = blocks._Residual(
filters=16,
strides=(1, 1),
input_shape=input_shape,
data_format="channels_last")
dx_tr, grads_tr, vars_tr = residual.backward_grads_and_vars(
x, dy=dy, training=True)
dx_ev, grads_ev, vars_ev = residual.backward_grads_and_vars(
x, dy=dy, training=False)
self.assertNotAllClose(dx_tr, dx_ev)
self.assertTrue(isinstance(grads_tr, list))
self.assertTrue(isinstance(grads_ev, list))
self.assertTrue(isinstance(vars_tr, list))
self.assertTrue(isinstance(vars_ev, list))
for grad_tr, var_tr, grad_ev, var_ev in zip(grads_tr, vars_tr, grads_ev,
vars_ev):
if grad_tr is not None: # Batch norm moving mean, var gives None grad
self.assertEqual(grad_tr.shape, grad_ev.shape)
self.assertEqual(var_tr.shape, var_ev.shape)
self.assertEqual(grad_tr.shape, var_tr.shape)
def test_backward_grads_and_vars_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16,) + input_shape
x = tf.random_normal(shape=data_shape)
dy = tf.random_normal(shape=data_shape)
residual = blocks._Residual(
filters=16, strides=(1, 1), input_shape=input_shape)
dx_tr, grads_tr, vars_tr = residual.backward_grads_and_vars(
x, dy=dy, training=True)
dx_ev, grads_ev, vars_ev = residual.backward_grads_and_vars(
x, dy=dy, training=False)
self.assertNotAllClose(dx_tr, dx_ev)
self.assertTrue(isinstance(grads_tr, list))
self.assertTrue(isinstance(grads_ev, list))
self.assertTrue(isinstance(vars_tr, list))
self.assertTrue(isinstance(vars_ev, list))
for grad_tr, var_tr, grad_ev, var_ev in zip(grads_tr, vars_tr, grads_ev,
vars_ev):
if grad_tr is not None: # Batch norm moving mean, var gives None grad
self.assertEqual(grad_tr.shape, grad_ev.shape)
self.assertEqual(var_tr.shape, var_ev.shape)
self.assertEqual(grad_tr.shape, var_tr.shape)
def test_backward_channels_last(self):
"""Test `backward` function with `channels_last` data format."""
with tf.device("/cpu:0"): # NHWC format
input_shape = (224, 224, 16)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape,
data_format="channels_last")
y_tr, y_ev = residual(x, training=True), residual(x,
training=False)
x_ = residual.backward(y_tr, training=True)
# Egregious numerical error
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
x_ = residual.backward(y_ev, training=False)
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
def test_backward_channels_last(self):
"""Test `backward` function with `channels_last` data format."""
with tf.device("/cpu:0"): # NHWC format
input_shape = (224, 224, 16)
data_shape = (16,) + input_shape
x = tf.random_normal(shape=data_shape)
residual = blocks._Residual(
filters=16,
strides=(1, 1),
input_shape=input_shape,
data_format="channels_last")
y_tr, y_ev = residual(x, training=True), residual(x, training=False)
x_ = residual.backward(y_tr, training=True)
# Egregious numerical error
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
x_ = residual.backward(y_ev, training=False)
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
def test_backward_channels_first(self):
"""Test `backward` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16,) + input_shape
x = tf.random_normal(shape=data_shape)
residual = blocks._Residual(
filters=16, strides=(1, 1), input_shape=input_shape)
y_tr, y_ev = residual(x, training=True), residual(x, training=False)
x_ = residual.backward(y_tr, training=True)
# The numerical loss is alarming; reconstructed inputs could differ from
# the original inputs often by more than 1e-3
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
x_ = residual.backward(y_ev, training=False)
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
def test_backward_grads_and_vars_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
dy = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape)
vars_and_vals = residual.get_moving_stats()
dx_tr, grads_tr, vars_tr = residual.backward_grads_and_vars(
x, dy=dy, training=True)
dx_ev, grads_ev, vars_ev = residual.backward_grads_and_vars(
x, dy=dy, training=False)
self.assertNotAllClose(dx_tr, dx_ev)
self.assertTrue(isinstance(grads_tr, list))
self.assertTrue(isinstance(grads_ev, list))
self.assertTrue(isinstance(vars_tr, list))
self.assertTrue(isinstance(vars_ev, list))
for grad_tr, var_tr, grad_ev, var_ev in zip(
grads_tr, vars_tr, grads_ev, vars_ev):
self.assertEqual(grad_tr.shape, grad_ev.shape)
self.assertEqual(var_tr.shape, var_ev.shape)
self.assertEqual(grad_tr.shape, var_tr.shape)
# Compare against the true gradient computed by the tape
residual.restore_moving_stats(vars_and_vals)
with tf.GradientTape(persistent=True) as tape:
tape.watch(x)
y = residual(x, training=True)
grads = tape.gradient(y, [x] + residual.trainable_variables,
output_gradients=[dy])
dx_tr_true, grads_tr_true = grads[0], grads[1:]
del tape
self.assertAllClose(dx_tr, dx_tr_true, rtol=1e-1, atol=1e-1)
self.assertAllClose(grads_tr, grads_tr_true, rtol=1e-1, atol=1e-1)
def test_backward_channels_first(self):
"""Test `backward` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape)
y_tr, y_ev = residual(x, training=True), residual(x,
training=False)
x_ = residual.backward(y_tr, training=True)
# The numerical loss is alarming; reconstructed inputs could differ from
# the original inputs often by more than 1e-3
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
x_ = residual.backward(y_ev, training=False)
self.assertAllClose(x, x_, rtol=1e-01, atol=1e-01)
def test_backward_channels_first(self):
"""Test `backward` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (16, 224, 224)
data_shape = (16, ) + input_shape
x = tf.random_normal(shape=data_shape)
residual = blocks._Residual(filters=16,
strides=(1, 1),
input_shape=input_shape)
y_tr, y_ev = residual(x, training=True), residual(x,
training=False)
x_ = residual.backward(y_ev, training=False)
self.assertAllClose(x, x_, rtol=1e-1, atol=1e-1)
x_ = residual.backward(y_tr,
training=True) # This updates moving avg
self.assertAllClose(x, x_, rtol=1e-1, atol=1e-1)
def test_backward_grads_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (128, 8, 8)
data_shape = (16, ) + input_shape
# Use double precision for testing
x_true = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy1, dy2 = tf.split(dy, num_or_size_splits=2, axis=1)
residual = blocks._Residual(filters=128,
strides=(1, 1),
input_shape=input_shape,
fused=False,
dtype=tf.float64)
with tf.GradientTape() as tape:
tape.watch(x_true)
x1_true, x2_true = tf.split(x_true,
num_or_size_splits=2,
axis=1)
y1, y2 = residual((x1_true, x2_true), training=True)
y = tf.concat((y1, y2), axis=1)
# Gradients computed due to reversibility
(x1, x2), (dx1, dx2), dw = residual.backward_grads(y=(y1, y2),
dy=(dy1, dy2),
training=True)
x = tf.concat((x1, x2), axis=1)
dx = tf.concat((dx1, dx2), axis=1)
# True gradients computed by the tape
grads = tape.gradient(y, [x_true] + residual.trainable_variables,
output_gradients=dy)
dx_true, dw_true = grads[0], grads[1:]
self.assertAllClose(x_true, x)
self.assertAllClose(dx_true, dx)
self.assertAllClose(dw_true, dw)
def test_backward_grads_channels_first(self):
"""Test `backward_grads` function with `channels_first` data format."""
if not tf.test.is_gpu_available():
self.skipTest("GPU not available")
with tf.device("/gpu:0"): # Default NCHW format
input_shape = (128, 8, 8)
data_shape = (16,) + input_shape
# Use double precision for testing
x_true = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy = tf.random_normal(shape=data_shape, dtype=tf.float64)
dy1, dy2 = tf.split(dy, num_or_size_splits=2, axis=1)
residual = blocks._Residual(
filters=128,
strides=(1, 1),
input_shape=input_shape,
fused=False,
dtype=tf.float64)
with tf.GradientTape() as tape:
tape.watch(x_true)
x1_true, x2_true = tf.split(x_true, num_or_size_splits=2, axis=1)
y1, y2 = residual((x1_true, x2_true), training=True)
y = tf.concat((y1, y2), axis=1)
# Gradients computed due to reversibility
(x1, x2), (dx1, dx2), dw = residual.backward_grads(
y=(y1, y2), dy=(dy1, dy2), training=True)
x = tf.concat((x1, x2), axis=1)
dx = tf.concat((dx1, dx2), axis=1)
# True gradients computed by the tape
grads = tape.gradient(
y, [x_true] + residual.trainable_variables, output_gradients=dy)
dx_true, dw_true = grads[0], grads[1:]
self.assertAllClose(x_true, x)
self.assertAllClose(dx_true, dx)
self.assertAllClose(dw_true, dw)