def check_x(layer):
y = layer(x)
y_mean, y_var = T.calculate_mean_and_var(
y, axis=[-T.rank(x)] + get_spatial_axis(spatial_ndims))
if use_bias:
assert_allclose(y_mean,
T.zeros_like(y_mean),
atol=1e-6,
rtol=1e-4)
assert_allclose(y_var,
T.ones_like(y_var),
atol=1e-6,
rtol=1e-4)
def check_x(layer):
y = layer(x)
y_mean, y_var = T.calculate_mean_and_var(y,
axis=T.int_range(
-T.rank(x),
-1))
if use_bias:
assert_allclose(y_mean,
T.zeros_like(y_mean),
atol=1e-6,
rtol=1e-4)
assert_allclose(y_var,
T.ones_like(y_var),
atol=1e-6,
rtol=1e-4)
def test_batch_norm_full_init(self):
T.random.seed(1234)
orig_input = T.random.randn([35, 4])
inputs = T.split(orig_input, sections=[7] * 5, axis=0)
mean, var = T.calculate_mean_and_var(orig_input, axis=[0])
layer = tk.layers.Sequential([
tk.layers.BatchNorm(4, momentum=0.125),
])
def step_fn(input):
_ = layer(input)
def data_generator():
for input in inputs:
yield (input, )
for i in range(4):
if i == 2:
loop = mltk.TrainLoop()
elif i == 3:
loop = mltk.TestLoop()
else:
loop = None
fn = lambda: tk.layers.batch_norm_full_init(
layer, data_generator(), step_fn, loop=loop)
tk.layers.set_train_mode(layer)
if loop is not None:
with loop:
fn()
else:
fn()
tk.layers.set_eval_mode(layer)
if T.backend_name == 'PyTorch':
self.assertEqual(layer[0].momentum, 0.125)
assert_allclose(layer[0].running_mean,
mean,
atol=1e-4,
rtol=1e-6)
else:
raise NotImplementedError()
def do_check(batch_shape, scale_type, initialized, dtype):
x = T.random.randn(make_conv_shape(batch_shape, num_features,
[6, 7, 8][:spatial_ndims]),
dtype=dtype)
# check construct
flow = cls(num_features,
scale=scale_type,
initialized=initialized,
dtype=dtype)
ctx.assertIn(f'num_features={num_features}', repr(flow))
ctx.assertIn(f'axis={-(spatial_ndims + 1)}', repr(flow))
ctx.assertIn(f'scale_type={scale_type!r}', repr(flow))
flow = tk.layers.jit_compile(flow)
# check initialize
if not initialized:
# must initialize with sufficient data
with pytest.raises(Exception, match='with at least .* dimensions'):
_ = flow(
T.random.randn(make_conv_shape([], num_features,
[6, 7, 8][:spatial_ndims]),
dtype=dtype))
# must initialize with inverse = Fale
with pytest.raises(Exception,
match='`ActNorm` must be initialized with '
'`inverse = False`'):
_ = flow(x, inverse=True)
# do initialize
y, _ = flow(x, compute_log_det=False)
y_mean, y_var = T.calculate_mean_and_var(
y, axis=[a for a in range(-T.rank(y), 0) if a != channel_axis])
assert_allclose(y_mean,
T.zeros([num_features]),
rtol=1e-4,
atol=1e-6)
assert_allclose(y_var,
T.ones([num_features]),
rtol=1e-4,
atol=1e-6)
else:
y, _ = flow(x, compute_log_det=False)
assert_allclose(y, x, rtol=1e-4, atol=1e-6)
# prepare for the expected result
scale_obj = ExpScale() if scale_type == 'exp' else LinearScale()
if T.IS_CHANNEL_LAST:
aligned_shape = [num_features]
else:
aligned_shape = [num_features] + [1] * spatial_ndims
bias = T.reshape(flow.bias, aligned_shape)
pre_scale = T.reshape(flow.pre_scale, aligned_shape)
expected_y, expected_log_det = scale_obj(x + bias,
pre_scale,
event_ndims=(spatial_ndims +
1),
compute_log_det=True)
flow_standard_check(ctx, flow, x, expected_y, expected_log_det,
T.random.randn(T.shape(expected_log_det)))
def test_data_dependent_initializer(self):
data_init = _MyDataDependentInitializer([])
# construct with the initializer
data_init.watcher.clear()
layer = tk.layers.Sequential(
tk.layers.Linear(5, 3, data_init=data_init))
self.assertEqual(data_init.watcher, [])
x = T.random.randn([2, 5])
y = T.random.randn([2, 5])
_ = layer(x)
_ = layer(y)
self.assertListEqual(data_init.watcher, [(layer[0], [x])])
# set_initialized(False) to re-enable the initializer
data_init.watcher.clear()
tk.init.set_initialized(layer, False)
x = T.random.randn([2, 5])
y = T.random.randn([2, 5])
_ = layer(x)
_ = layer(y)
self.assertListEqual(data_init.watcher, [(layer[0], [x])])
# set_initialize(True) to disable newly constructed data-init
data_init.watcher.clear()
layer = tk.layers.Sequential(
tk.layers.Linear(5, 3, data_init=data_init))
tk.init.set_initialized(layer, True)
x = T.random.randn([2, 5])
_ = layer(x)
self.assertListEqual(data_init.watcher, [])
# remove the data-dependent initializers
data_init.watcher.clear()
layer = tk.layers.Sequential(
tk.layers.Linear(5, 3, data_init=data_init))
tk.init.remove_data_dependent_initializers(layer)
tk.init.set_initialized(layer, False)
x = T.random.randn([2, 5])
_ = layer(x)
self.assertListEqual(data_init.watcher, [])
# also `set_initialized` will affect layers with `set_initialized()`
# method, e.g., `ActNorm`
x = T.random.randn([2, 3, 5])
layer = tk.layers.jit_compile(tk.layers.ActNorm(5))
self.assertFalse(layer.flow.initialized)
tk.init.set_initialized(layer)
self.assertTrue(layer.flow.initialized)
assert_allclose(layer(x), x, rtol=1e-4, atol=1e-6)
tk.init.set_initialized(layer, False)
self.assertFalse(layer.flow.initialized)
y = layer(x)
y_mean, y_var = T.calculate_mean_and_var(y, axis=[0, 1])
assert_allclose(y_mean, T.zeros_like(y_mean), rtol=1e-4, atol=1e-6)
assert_allclose(y_var, T.ones_like(y_var), rtol=1e-4, atol=1e-6)
self.assertTrue(layer.flow.initialized)
assert_allclose(layer(x), y, rtol=1e-4, atol=1e-6)