def test_sequential_with_quadratic_flows(self):
n_layers = 3
flow1 = MultiLayerQuadraticFlow(n_layers)
flow2 = SequentialFlow([
QuadraticFlow(i + 1., i * 2. + 1.)
for i in range(n_layers)
])
self.assertTrue(flow2.explicitly_invertible)
self.assertEqual(len(flow2.flows), n_layers)
for i in range(n_layers):
self.assertEqual(flow2.flows[i].a, i + 1.)
self.assertEqual(flow2.flows[i].b, i * 2. + 1.)
x = tf.range(12, dtype=tf.float32) + 1.
with self.test_session() as sess:
invertible_flow_standard_check(self, flow2, sess, x)
# transform
y1, log_det_y1 = flow1.transform(x)
y2, log_det_y2 = flow2.transform(x)
np.testing.assert_allclose(*sess.run([y1, y2]))
np.testing.assert_allclose(*sess.run([log_det_y1, log_det_y2]))
# inverse transform
x1, log_det_x1 = flow1.inverse_transform(y1)
x2, log_det_x2 = flow1.inverse_transform(y2)
np.testing.assert_allclose(*sess.run([x1, x2]))
np.testing.assert_allclose(*sess.run([log_det_x1, log_det_x2]))
def test_invert_flow(self):
with self.test_session() as sess:
# test invert a normal flow
flow = QuadraticFlow(2., 5.)
inv_flow = flow.invert()
self.assertIsInstance(inv_flow, InvertFlow)
self.assertEqual(inv_flow.x_value_ndims, 0)
self.assertEqual(inv_flow.y_value_ndims, 0)
self.assertFalse(inv_flow.require_batch_dims)
test_x = np.arange(12, dtype=np.float32) + 1.
test_y, test_log_det = quadratic_transform(npyops, test_x, 2., 5.)
self.assertFalse(flow._has_built)
y, log_det_y = inv_flow.inverse_transform(tf.constant(test_x))
self.assertTrue(flow._has_built)
np.testing.assert_allclose(sess.run(y), test_y)
np.testing.assert_allclose(sess.run(log_det_y), test_log_det)
invertible_flow_standard_check(self, inv_flow, sess, test_y)
# test invert an InvertFlow
inv_inv_flow = inv_flow.invert()
self.assertIs(inv_inv_flow, flow)
# test use with FlowDistribution
normal = Normal(mean=1., std=2.)
inv_flow = QuadraticFlow(2., 5.).invert()
distrib = FlowDistribution(normal, inv_flow)
distrib_log_det = distrib.log_prob(test_x)
np.testing.assert_allclose(*sess.run(
[distrib_log_det,
normal.log_prob(test_y) + test_log_det]))
def test_feature_shuffling_flow(self):
np.random.seed(1234)
with self.test_session() as sess:
# axis = -1, value_ndims = 1
x = np.random.normal(size=[3, 4, 5, 6]).astype(np.float32)
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 6])
permutation = np.arange(6, dtype=np.int32)
np.random.shuffle(permutation)
y = x[..., permutation]
log_det = np.zeros([3, 4, 5]).astype(np.float32)
layer = FeatureShufflingFlow(axis=-1, value_ndims=1)
y_out, log_det_out = layer.transform(x_ph)
sess.run(tf.assign(layer._permutation, permutation))
y_out, log_det_out = sess.run(
[y_out, log_det_out], feed_dict={x_ph: x})
np.testing.assert_equal(y_out, y)
np.testing.assert_equal(log_det_out, log_det)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
assert_variables(['permutation'], trainable=False,
scope='feature_shuffling_flow',
collections=[tf.GraphKeys.MODEL_VARIABLES])
# axis = -2, value_ndims = 3
x = np.random.normal(size=[3, 4, 5, 6]).astype(np.float32)
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 5, None])
permutation = np.arange(5, dtype=np.int32)
np.random.shuffle(permutation)
y = x[..., permutation, :]
log_det = np.zeros([3]).astype(np.float32)
layer = FeatureShufflingFlow(axis=-2, value_ndims=3)
y_out, log_det_out = layer.transform(x_ph)
sess.run(tf.assign(layer._permutation, permutation))
y_out, log_det_out = sess.run(
[y_out, log_det_out], feed_dict={x_ph: x})
np.testing.assert_equal(y_out, y)
np.testing.assert_equal(log_det_out, log_det)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
def test_invertible_dense(self):
assert_allclose = functools.partial(np.testing.assert_allclose,
rtol=1e-5)
np.random.seed(1234)
with self.test_session() as sess:
x = np.random.normal(size=[3, 5, 7]).astype(np.float32)
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 7])
kernel = np.random.normal(size=(7, 7)).astype(x.dtype)
y, log_det = naive_invertible_linear(x, kernel, -1, 1)
layer = InvertibleDense(strict_invertible=False)
y_out, log_det_out = layer.transform(x_ph)
# The kernel is initialized as an orthogonal matrix. As a result,
# the log-det is zero (the computation result should be close to
# zero, but may not be zero). This is not good for testing.
# Thus we initialize it with an arbitrary matrix.
sess.run(tf.assign(layer._kernel_matrix._matrix, kernel))
y_out, log_det_out = sess.run([y_out, log_det_out],
feed_dict={x_ph: x})
assert_allclose(y_out, y)
assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self,
layer,
sess,
x_ph,
feed_dict={x_ph: x},
rtol=1e-5,
atol=1e-6)
assert_variables(['matrix'],
trainable=True,
scope='invertible_dense/kernel',
collections=[tf.GraphKeys.MODEL_VARIABLES])
# test non-trainable
with tf.Graph().as_default():
layer = InvertibleDense(strict_invertible=False, trainable=False)
layer.apply(tf.zeros([2, 3, 4, 5]))
assert_variables(['matrix'],
trainable=False,
scope='invertible_dense/kernel',
collections=[tf.GraphKeys.MODEL_VARIABLES])
def check(x, x_value_ndims, y_value_shape, x_ph=None):
flow = ReshapeFlow(x_value_ndims=x_value_ndims,
y_value_shape=y_value_shape)
if x_value_ndims > 0:
batch_shape = list(x.shape[:-x_value_ndims])
else:
batch_shape = list(x.shape)
y_ans = np.reshape(x, batch_shape + list(y_value_shape))
log_det_ans = np.zeros(batch_shape)
feed_dict = None
if x_ph is not None:
feed_dict = {x_ph: x}
x = x_ph
y, log_det = sess.run(flow.transform(x), feed_dict=feed_dict)
np.testing.assert_allclose(y, y_ans)
np.testing.assert_allclose(log_det, log_det_ans)
invertible_flow_standard_check(self, flow, sess, x, feed_dict)
def check(x, y_shape, bs, channels_last=True, x_ph=None):
# compute the answer
y = naive_space_to_depth(x, bs, channels_last)
log_det = np.zeros(y.shape[:-3], dtype=np.float32)
self.assertTupleEqual(y.shape, y_shape)
# get feed dict
feed_dict = None
if x_ph is not None:
feed_dict = {x_ph: x}
x = x_ph
flow = SpaceToDepthFlow(bs, channels_last=channels_last)
y_out, log_det_out = sess.run(flow.transform(x),
feed_dict=feed_dict)
np.testing.assert_allclose(y_out, y)
np.testing.assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self, flow, sess, x, feed_dict)
def test_different_value_ndims(self):
def reshape_tail(x, value_ndims, shape):
batch_shape = x.shape
if value_ndims > 0:
batch_shape = batch_shape[:-value_ndims]
return np.reshape(x, batch_shape + tuple(shape))
def split_transform(x,
split_axis,
join_axis,
x_value_ndims,
y_shape,
a1,
b1,
a2=None,
b2=None):
n1 = x.shape[split_axis] // 2
n2 = x.shape[split_axis] - n1
x1, x2 = np.split(x, [n1], axis=split_axis)
y1, log_det1 = quadratic_transform(npyops, x1, a1, b1)
if a2 is not None:
y2, log_det2 = quadratic_transform(npyops, x2, a2, b2)
else:
y2, log_det2 = x2, np.zeros_like(x, dtype=np.float64)
y1 = reshape_tail(y1, x_value_ndims, y_shape)
y2 = reshape_tail(y2, x_value_ndims, y_shape)
y = np.concatenate([y1, y2], axis=join_axis)
if x_value_ndims > 0:
reduce_axis = tuple(range(-x_value_ndims, 0))
log_det1 = np.sum(log_det1, axis=reduce_axis)
log_det2 = np.sum(log_det2, axis=reduce_axis)
log_det = log_det1 + log_det2
return y, log_det
with self.test_session() as sess:
np.random.seed(1234)
x = 10. * np.random.normal(size=[3, 4, 5, 12]).astype(np.float64)
# 2 -> 3, x_value_ndims = 3, y_value_ndims = 4
x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4)
flow = SplitFlow(split_axis=-2,
join_axis=2,
left=SequentialFlow([
QuadraticFlow(2., 5., value_ndims=3),
ReshapeFlow(3, [4, -1, 2, 6]),
]),
right=SequentialFlow([
QuadraticFlow(1.5, 3., value_ndims=3),
ReshapeFlow(3, [4, -1, 2, 6]),
]))
self.assertEqual(flow.x_value_ndims, 3)
self.assertEqual(flow.y_value_ndims, 4)
y, log_det = split_transform(x,
split_axis=-2,
join_axis=-3,
x_value_ndims=3,
y_shape=[4, -1, 2, 6],
a1=2.,
b1=5.,
a2=1.5,
b2=3.)
y_out, log_det_out = sess.run(flow.transform(x_ph),
feed_dict={x_ph: x})
np.testing.assert_allclose(y_out, y)
np.testing.assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self,
flow,
sess,
x_ph,
feed_dict={x_ph: x})
def test_equal_value_ndims(self):
def split_transform(x,
split_axis,
value_ndims,
a1,
b1,
a2=None,
b2=None):
n1 = x.shape[split_axis] // 2
n2 = x.shape[split_axis] - n1
x1, x2 = np.split(x, [n1], axis=split_axis)
y1, log_det1 = quadratic_transform(npyops, x1, a1, b1)
if a2 is not None:
y2, log_det2 = quadratic_transform(npyops, x2, a2, b2)
else:
y2, log_det2 = x2, np.zeros_like(x, dtype=np.float64)
y = np.concatenate([y1, y2], axis=split_axis)
if value_ndims > 0:
reduce_axis = tuple(range(-value_ndims, 0))
log_det1 = np.sum(log_det1, axis=reduce_axis)
log_det2 = np.sum(log_det2, axis=reduce_axis)
log_det = log_det1 + log_det2
return y, log_det
with self.test_session() as sess:
np.random.seed(1234)
x = 10. * np.random.normal(size=[3, 4, 5, 6]).astype(np.float64)
# static input, split_axis = -1, value_ndims = 1, right = None
flow = SplitFlow(-1, QuadraticFlow(2., 5., value_ndims=1))
self.assertEqual(flow.x_value_ndims, 1)
self.assertEqual(flow.y_value_ndims, 1)
y, log_det = split_transform(x,
split_axis=-1,
value_ndims=1,
a1=2.,
b1=5.)
y_out, log_det_out = sess.run(flow.transform(x))
np.testing.assert_allclose(y_out, y)
np.testing.assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self,
flow,
sess,
x,
rtol=1e-4,
atol=1e-5)
# dynamic input, split_axis = -2, value_ndims = 2, right = None
x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4)
flow = SplitFlow(-2, QuadraticFlow(2., 5., value_ndims=2))
self.assertEqual(flow.x_value_ndims, 2)
self.assertEqual(flow.y_value_ndims, 2)
y, log_det = split_transform(x,
split_axis=-2,
value_ndims=2,
a1=2.,
b1=5.)
y_out, log_det_out = sess.run(flow.transform(x_ph),
feed_dict={x_ph: x})
np.testing.assert_allclose(y_out, y)
np.testing.assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self,
flow,
sess,
x_ph,
feed_dict={x_ph: x})
# dynamic input, split_axis = 2, value_ndims = 3
x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4)
flow = SplitFlow(split_axis=2,
left=QuadraticFlow(2., 5., value_ndims=3),
right=QuadraticFlow(1.5, 3., value_ndims=3))
self.assertEqual(flow.x_value_ndims, 3)
self.assertEqual(flow.y_value_ndims, 3)
y, log_det = split_transform(x,
split_axis=2,
value_ndims=3,
a1=2.,
b1=5.,
a2=1.5,
b2=3.)
y_out, log_det_out = sess.run(flow.transform(x_ph),
feed_dict={x_ph: x})
np.testing.assert_allclose(y_out, y)
np.testing.assert_allclose(log_det_out, log_det)
invertible_flow_standard_check(self,
flow,
sess,
x_ph,
feed_dict={x_ph: x})
def test_act_norm(self):
assert_allclose = functools.partial(
np.testing.assert_allclose, rtol=1e-5, atol=1e-6)
np.random.seed(1234)
x = np.random.normal(size=[3, 4, 5, 6, 7])
x_ph = tf.placeholder(dtype=tf.float64, shape=[None, None, 5, None, 7])
x2 = np.random.normal(size=[2, 3, 4, 5, 6, 7])
x2_ph = tf.placeholder(dtype=tf.float64,
shape=[None, None, None, 5, None, 7])
x3 = np.random.normal(size=[4, 5, 6, 7])
x3_ph = tf.placeholder(dtype=tf.float64, shape=[None, 5, None, 7])
with self.test_session() as sess:
# -- static input shape, scale_type = 'linear', value_ndims = 3
axis = [-1, -3]
value_ndims = 3
var_shape = (5, 7)
scale, bias, var_shape_aligned = naive_act_norm_initialize(x, axis)
self.assertEqual(scale.shape, var_shape)
self.assertEqual(bias.shape, var_shape)
# test initialize
act_norm = ActNorm(axis=axis, value_ndims=value_ndims,
scale_type='linear')
y_out, log_det_out = sess.run(
act_norm.transform(tf.constant(x, dtype=tf.float64)))
self.assertEqual(act_norm._bias.dtype.base_dtype, tf.float64)
scale_out, bias_out = sess.run(
[act_norm._pre_scale, act_norm._bias])
assert_allclose(scale_out, scale)
assert_allclose(bias_out, bias)
# test the transform output from the initializing procedure
y, log_det = naive_act_norm_transform(
x, var_shape_aligned, value_ndims, scale, bias)
self.assertEqual(y.shape, x.shape)
self.assertEqual(log_det.shape, x.shape[:-value_ndims])
assert_allclose(y_out, y)
assert_allclose(log_det_out, log_det)
# test use an initialized act_norm
y2, log_det2 = naive_act_norm_transform(
x2, var_shape_aligned, value_ndims, scale, bias)
self.assertEqual(y2.shape, x2.shape)
self.assertEqual(log_det2.shape, x2.shape[:-value_ndims])
y2_out, log_det2_out = sess.run(
act_norm.transform(x2_ph), feed_dict={x2_ph: x2})
assert_allclose(y2_out, y2)
assert_allclose(log_det2_out, log_det2)
# invertible flow standard checks
invertible_flow_standard_check(self, act_norm, sess, x)
invertible_flow_standard_check(
self, act_norm, sess, x2_ph, feed_dict={x2_ph: x2})
invertible_flow_standard_check(
self, act_norm, sess, x3_ph, feed_dict={x3_ph: x3})
# -- dynamic input shape, scale_type = 'exp', value_ndims = 4
value_ndims = 4
# test initialize
act_norm = ActNorm(axis=axis, value_ndims=value_ndims,
scale_type='exp')
y_out, log_det_out = sess.run(
act_norm.transform(x_ph), feed_dict={x_ph: x})
self.assertEqual(act_norm._bias.dtype.base_dtype, tf.float64)
scale_out, bias_out = sess.run(
[tf.exp(act_norm._pre_scale), act_norm._bias])
assert_allclose(scale_out, scale)
assert_allclose(bias_out, bias)
# test the transform output from the initializing procedure
y, log_det = naive_act_norm_transform(
x, var_shape_aligned, value_ndims, scale, bias)
self.assertEqual(y.shape, x.shape)
self.assertEqual(log_det.shape, x.shape[:-value_ndims])
assert_allclose(y_out, y)
assert_allclose(log_det_out, log_det)
# test use an initialized act_norm
y3, log_det3 = naive_act_norm_transform(
x3, var_shape_aligned, value_ndims, scale, bias)
self.assertEqual(y3.shape, x3.shape)
self.assertEqual(log_det3.shape, x3.shape[:-value_ndims])
y3_out, log_det3_out = sess.run(
act_norm.transform(x3_ph), feed_dict={x3_ph: x3})
assert_allclose(y3_out, y3)
assert_allclose(log_det3_out, log_det3)
# invertible flow standard checks
invertible_flow_standard_check(self, act_norm, sess, x)
invertible_flow_standard_check(
self, act_norm, sess, x2_ph, feed_dict={x2_ph: x2})
invertible_flow_standard_check(
self, act_norm, sess, x3_ph, feed_dict={x3_ph: x3})
def test_coupling_layer(self):
assert_allclose = functools.partial(
np.testing.assert_allclose, rtol=1e-5)
np.random.seed(1234)
kernel1 = np.random.normal(size=[3, 2]).astype(np.float32)
kernel2 = np.random.normal(size=[2, 3]).astype(np.float32)
shift1 = np.random.normal(size=[2]).astype(np.float32)
shift2 = np.random.normal(size=[3]).astype(np.float32)
def shift_and_scale_fn(x1, n2, no_scale=False):
kernel = kernel1 if n2 == 2 else kernel2
shift = tf.convert_to_tensor(shift1 if n2 == 2 else shift2)
assert(kernel.shape[-1] == n2)
assert(shift.shape[-1] == n2)
x1, s1, s2 = flatten_to_ndims(x1, 2)
scale = unflatten_from_ndims(tf.matmul(x1, kernel), s1, s2)
shift = shift + tf.zeros_like(scale, dtype=shift.dtype)
if no_scale:
scale = None
return shift, scale
def shift_and_scale_numpy_fn(x1, n2, no_scale=False):
a, b = shift_and_scale_fn(x1, n2, no_scale=no_scale)
if b is None:
a = sess.run(a)
else:
a, b = sess.run([a, b])
return a, b
with self.test_session() as sess:
# test linear scale, primary
x = np.random.normal(size=[3, 4, 5]).astype(np.float32)
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 5])
axis = -1
value_ndims = 1
y_ans, log_det_ans = naive_coupling_layer(
shift_and_scale_numpy_fn, x, axis=axis,
value_ndims=value_ndims, secondary=False, scale_type='linear',
reverse=False
)
layer = CouplingLayer(
shift_and_scale_fn, axis=axis, value_ndims=value_ndims,
secondary=False, scale_type='linear'
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
# test exp scale, primary
axis = -1
value_ndims = 2
y_ans, log_det_ans = naive_coupling_layer(
shift_and_scale_numpy_fn, x, axis=axis,
value_ndims=value_ndims, secondary=False, scale_type='exp',
reverse=False
)
layer = CouplingLayer(
shift_and_scale_fn, axis=axis, value_ndims=value_ndims,
secondary=False, scale_type='exp'
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
# test sigmoid scale, secondary
sigmoid_scale_bias = np.exp(1)
axis = -1
value_ndims = 1
y_ans, log_det_ans = naive_coupling_layer(
shift_and_scale_numpy_fn, x, axis=axis,
value_ndims=value_ndims, secondary=False, scale_type='sigmoid',
sigmoid_scale_bias=sigmoid_scale_bias, reverse=False
)
layer = CouplingLayer(
shift_and_scale_fn, axis=axis, value_ndims=value_ndims,
secondary=False, scale_type='sigmoid',
sigmoid_scale_bias=sigmoid_scale_bias
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
# test None scale, primary
axis = -1
value_ndims = 1
y_ans, log_det_ans = naive_coupling_layer(
functools.partial(shift_and_scale_numpy_fn, no_scale=True),
x, axis=axis,
value_ndims=value_ndims, secondary=False, scale_type=None,
reverse=False
)
layer = CouplingLayer(
functools.partial(shift_and_scale_fn, no_scale=True),
axis=axis, value_ndims=value_ndims,
secondary=False, scale_type=None
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
# test None scale, secondary
axis = -1
value_ndims = 3
y_ans, log_det_ans = naive_coupling_layer(
functools.partial(shift_and_scale_numpy_fn, no_scale=True),
x, axis=axis,
value_ndims=value_ndims, secondary=True, scale_type=None,
reverse=False
)
layer = CouplingLayer(
functools.partial(shift_and_scale_fn, no_scale=True),
axis=axis, value_ndims=value_ndims,
secondary=True, scale_type=None
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x})
def test_coupling_layer_with_conv2d(self):
assert_allclose = functools.partial(
np.testing.assert_allclose, atol=1e-5, rtol=5e-4)
np.random.seed(1234)
kernel1 = np.random.normal(size=[3, 3, 5, 6]).astype(np.float32)
kernel2 = np.random.normal(size=[3, 3, 6, 5]).astype(np.float32)
shift1 = np.random.normal(size=[6]).astype(np.float32)
shift2 = np.random.normal(size=[5]).astype(np.float32)
def shift_and_scale_fn(x1, n2, no_scale=False, channels_last=True):
kernel = kernel1 if n2 == 6 else kernel2
shift = tf.convert_to_tensor(shift1 if n2 == 6 else shift2)
assert (kernel.shape[-1] == n2)
assert (shift.shape[-1] == n2)
x1 = transpose_conv2d_channels_x_to_last(
x1, channels_last=channels_last
)
scale = conv2d(x1, n2, (3, 3), use_bias=False, kernel=kernel,
channels_last=True)
shift = shift + tf.zeros_like(scale, dtype=shift.dtype)
scale = transpose_conv2d_channels_last_to_x(scale, channels_last)
shift = transpose_conv2d_channels_last_to_x(shift, channels_last)
if no_scale:
scale = None
return shift, scale
def shift_and_scale_numpy_fn(x1, n2, no_scale=False,
channels_last=True):
a, b = shift_and_scale_fn(x1, n2, no_scale, channels_last)
if b is None:
a = sess.run(a)
else:
a, b = sess.run([a, b])
return a, b
with self.test_session() as sess:
# test exp scale, primary, NHWC
x = np.random.normal(size=[11, 13, 32, 31, 11]).astype(np.float32)
x_ph = tf.placeholder(dtype=tf.float32,
shape=[None, None, None, None, 11])
axis = -1
value_ndims = 3
y_ans, log_det_ans = naive_coupling_layer(
shift_and_scale_numpy_fn, x, axis=axis,
value_ndims=value_ndims, secondary=False, scale_type='exp',
reverse=False
)
layer = CouplingLayer(
shift_and_scale_fn, axis=axis, value_ndims=value_ndims,
secondary=False, scale_type='exp'
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x}, rtol=5e-4,
atol=1e-5
)
# test sigmoid scale, secondary, NCHW
x = np.transpose(x, [0, 1, 4, 2, 3])
x_ph = tf.placeholder(dtype=tf.float32,
shape=[None, None, 11, None, None])
axis = -3
value_ndims = 3
y_ans, log_det_ans = naive_coupling_layer(
functools.partial(shift_and_scale_numpy_fn,
channels_last=False),
x, axis=axis,
value_ndims=value_ndims, secondary=True, scale_type='sigmoid',
reverse=False
)
layer = CouplingLayer(
functools.partial(shift_and_scale_fn, channels_last=False),
axis=axis, value_ndims=value_ndims,
secondary=True, scale_type='sigmoid'
)
y, log_det = layer.transform(x_ph)
y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x})
assert_allclose(y_out, y_ans)
assert_allclose(log_det_out, log_det_ans)
invertible_flow_standard_check(
self, layer, sess, x_ph, feed_dict={x_ph: x}, rtol=5e-4,
atol=1e-5
)
