def test_actnorm_bias_init_conv(self):
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.ActnormBiasLayer()
with self.assertRaises(AssertionError):
layer.get_ddi_init_ops()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops()
self.assertEqual(z, None)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
sess.run(init_ops)
x, logdet = sess.run([x, logdet])
bias = sess.run(layer._bias_t)
self.assertEqual(x.shape, images_np.shape)
self.assertEqual(bias.shape, (1, 1, 1, 3))
self.assertGreater(np.sum(bias**2), 0)
# check mean after passing act norm
self.assertAllClose(np.mean(x.reshape([-1, 3]), axis=0),
[0.0, 0.0, 0.0])
self.forward_inverse(layer, flow)
def test_actnorm_scale_init_conv_iter(self):
np.random.seed(52321)
images_ph = tf.placeholder(tf.float32, shape=[8, 32, 32, 3])
flow = fl.InputLayer(images_ph)
layer = fl.ActnormScaleLayer(scale=np.sqrt(np.pi))
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops(num_init_iterations=50)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
for i in range(200):
sess.run(init_ops,
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
for i in range(5):
x_np, logdet_np = sess.run(
[x, logdet],
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
self.assertEqual(x.shape, x_np.shape)
self.assertAllClose(
np.var(x_np.reshape([-1, 3]), axis=0),
[np.pi, np.pi, np.pi],
atol=0.1,
)
self.forward_inverse(
layer, flow, feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
def test_actnorm_scale_init_conv(self):
np.random.seed(52321)
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.ActnormScaleLayer()
with self.assertRaises(AssertionError):
layer.get_ddi_init_ops()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops()
self.assertEqual(z, None)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
sess.run(init_ops)
x, logdet = sess.run([x, logdet])
log_scale = sess.run(layer._log_scale_t)
self.assertEqual(x.shape, images_np.shape)
self.assertEqual(log_scale.shape, (1, 1, 1, 3))
self.assertGreater(np.sum(log_scale**2), 0)
self.assertGreater(np.sum(logdet**2), 0)
# check var after passing act norm
self.assertAllClose(np.var(x.reshape([-1, 3]), axis=0),
[1.0, 1.0, 1.0],
atol=0.001)
self.forward_inverse(layer, flow)
def test_actnorm_init_conv(self):
np.random.seed(52321)
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.ActnormLayer(scale=np.sqrt(np.pi))
with self.assertRaises(AssertionError):
layer.get_ddi_init_ops()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops()
self.assertEqual(z, None)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
sess.run(init_ops)
x, logdet = sess.run([x, logdet])
self.assertEqual(x.shape, images_np.shape)
# check var after passing act norm
self.assertAllClose(np.var(x.reshape([-1, 3]), axis=0),
[np.pi] * 3,
atol=0.01)
self.assertAllClose(np.mean(x.reshape([-1, 3]), axis=0), [0.0] * 3,
atol=0.01)
self.forward_inverse(layer, flow)
def test_logitify_layer_conv(self):
np.random.seed(52321)
images_np = np.random.rand(8, 32, 32, 1)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.LogitifyImage()
self.forward_inverse(layer, flow, atol=0.01)
x, logdet, z = flow
logdet += 10.0
flow = x, logdet, z
new_flow = layer(flow, forward=True)
flow_rec = layer(new_flow, forward=False)
x, logdet, z = new_flow
x_rec, logdet_rec, z = flow_rec
self.assertEqual(z, None)
self.assertEqual(x.shape.as_list(), [8, 32, 32, 1])
self.assertEqual(logdet.shape.as_list(), [8])
self.assertEqual(x_rec.shape.as_list(), [8, 32, 32, 1])
self.assertEqual(logdet_rec.shape.as_list(), [8])
with self.test_session() as sess:
_ = sess.run([x, logdet])
x_rec, logdet_rec = sess.run([x_rec, logdet_rec])
self.assertAllClose(logdet_rec, [10.0] * 8, atol=0.01)
self.assertAllClose(images_np, x_rec, atol=0.01)
def test_actnorm_scale_conv(self):
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.ActnormScaleLayer()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
self.assertEqual(z, None)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
log_scale = sess.run(layer._log_scale_t)
x, logdet = sess.run([x, logdet])
self.assertEqual(x.shape, images_np.shape)
self.assertEqual(log_scale.shape, (1, 1, 1, 3))
self.assertEqual(np.sum(log_scale**2), 0)
# zero initialization
self.assertAllClose(logdet, [0.0] * 8)
self.forward_inverse(layer, flow)
def test_invertible_conv1x1_learn_identity(self):
images = tf.random_normal([8, 32, 32, 16])
flow = fl.InputLayer(images)
layer = fl.InvertibleConv1x1Layer(use_lu_decomposition=True)
self.try_to_train_identity_layer(layer, flow)
layer = fl.InvertibleConv1x1Layer(use_lu_decomposition=False)
self.try_to_train_identity_layer(layer, flow)
def test_input_layer(self):
images = np.random.rand(8, 32, 32, 1)
images = tf.to_float(images)
x, logdet, z = fl.InputLayer(images)
self.assertEqual(images, x)
self.assertEqual(logdet.shape.as_list(), [8])
self.assertEqual(z, None)
def encoder(self, x):
#if x is None:
# x = self.x_image_ph
#flow = fl.InputLayer(self.x_image_ph)
flow = fl.InputLayer(x)
output_flow = self.model_flow(flow, forward=True)
# ## Prepare output tensors
y, logdet, z = output_flow
return y, logdet, z
def test_squeezing_layer_conv(self):
images = np.random.rand(8, 32, 32, 1)
images = tf.to_float(images)
flow = fl.InputLayer(images)
layer = fl.SqueezingLayer()
self.forward_inverse(layer, flow)
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
self.assertEqual(logdet.shape.as_list(), [8])
self.assertEqual([8, 16, 16, 4], x.shape.as_list())
def test_quantize_image_layer_conv(self):
np.random.seed(52321)
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.QuantizeImage(num_bits=8)
self.forward_inverse(layer, flow, atol=1.5 / 256)
new_flow = layer(flow, forward=True)
flow_rec = layer(new_flow, forward=False)
x, logdet, z = new_flow
x_rec, logdet_rec, z = flow_rec
self.assertEqual(z, None)
self.assertEqual(x.shape.as_list(), [8, 32, 32, 3])
self.assertEqual(x_rec.shape.as_list(), [8, 32, 32, 3])
# less bits
flow = fl.InputLayer(images)
layer = fl.QuantizeImage(num_bits=5)
self.forward_inverse(layer, flow, atol=1.5 / 32)
layer = fl.QuantizeImage(num_bits=4)
new_flow = layer(flow, forward=True)
flow_rec = layer(new_flow, forward=False)
with self.test_session() as sess:
x_rec_uint8 = layer.to_uint8(flow_rec[0])
self.assertEqual(x_rec_uint8.dtype, tf.uint8)
x_rec_uint8 = sess.run(x_rec_uint8)
self.assertAllGreaterEqual(x_rec_uint8, 0)
self.assertAllLessEqual(x_rec_uint8, 255)
self.assertEqual(np.unique(x_rec_uint8).shape, (2**4, ))
with self.assertRaises(AssertionError):
layer = fl.QuantizeImage(num_bits=4)
self.forward_inverse(layer, flow, atol=1 / 32)
def test_actnorm_init_conv_iter(self):
np.random.seed(52321)
images_ph = tf.placeholder(tf.float32, shape=[8, 32, 32, 3])
flow = fl.InputLayer(images_ph)
layer = fl.ActnormLayer(scale=np.sqrt(np.pi))
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops(num_init_iterations=50)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
for i in range(200):
sess.run(init_ops,
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
for i in range(5):
x_np, logdet_np = sess.run(
[x, logdet],
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
self.assertEqual(x.shape, x_np.shape)
self.assertAllClose(np.var(x_np.reshape([-1, 3]), axis=0),
[np.pi] * 3,
atol=0.1)
self.assertAllClose(np.mean(x_np.reshape([-1, 3]), axis=0),
[0.0] * 3,
atol=0.1)
self.forward_inverse(
layer, flow, feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
def feed_dict_fn():
return {images_ph: np.random.rand(8, 32, 32, 3)}
def post_init_fn(sess):
init_ops = layer.get_ddi_init_ops()
sess.run(init_ops, {images_ph: np.random.rand(8, 32, 32, 3)})
with tf.variable_scope("TestTrain"):
layer = fl.ActnormLayer(scale=np.sqrt(np.pi))
self.try_to_train_identity_layer(layer,
flow,
feed_dict_fn=feed_dict_fn,
post_init_fn=post_init_fn)
def test_create_simple_flow(self):
np.random.seed(642201)
images = tf.placeholder(tf.float32, [16, 32, 32, 1])
layers, actnorm_layers = nets.create_simple_flow(num_steps=2,
num_scales=4,
num_bits=8)
flow = fl.InputLayer(images)
model_flow = fl.ChainLayer(layers)
output_flow = model_flow(flow, forward=True)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
for actnorm_layer in actnorm_layers:
init_op = actnorm_layer.get_ddi_init_ops(10)
noise = np.random.rand(16, 32, 32, 1)
# fit actnorms to certain noise
for i in range(30):
sess.run(init_op, feed_dict={images: noise})
actnorm_flow = actnorm_layer._forward_outputs[0]
normed_x = sess.run(actnorm_flow[0], feed_dict={images: noise})
nc = normed_x.shape[-1]
self.assertAllClose(np.var(normed_x.reshape([-1, nc]), axis=0),
[1.0] * nc,
atol=0.1)
self.assertAllClose(np.mean(normed_x.reshape([-1, nc]),
axis=0), [0.0] * nc,
atol=0.1)
output_flow_np = sess.run(
output_flow, feed_dict={images: np.random.rand(16, 32, 32, 1)})
y, logdet, z = output_flow_np
self.assertEqual(
np.prod(y.shape) + np.prod(z.shape), np.prod([16, 32, 32, 1]))
self.assertTrue(np.max(np.abs(y)) < 15.0)
self.forward_inverse(
sess,
model_flow,
flow,
atol=0.01,
feed_dict={images: np.random.rand(16, 32, 32, 1)},
)
def test_invertible_conv1x1_lu_decomp(self):
images_np = np.random.rand(8, 32, 32, 16)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.InvertibleConv1x1Layer(use_lu_decomposition=True)
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
self.assertEqual(z, None)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
x, logdet = sess.run([x, logdet])
self.assertEqual(x.shape, images_np.shape)
self.forward_inverse(layer, flow)
def test_actnorm_conv(self):
images_np = np.random.rand(8, 32, 32, 3)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.ActnormLayer()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
self.assertEqual(z, None)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
x, logdet = sess.run([x, logdet])
self.assertEqual(x.shape, images_np.shape)
self.forward_inverse(layer, flow)
def test_initialize_actnorms(self):
np.random.seed(642201)
images_ph = tf.placeholder(tf.float32, [16, 16, 16, 1])
layers, actnorm_layers = nets.create_simple_flow(num_steps=1,
num_scales=3)
flow = fl.InputLayer(images_ph)
model_flow = fl.ChainLayer(layers)
output_flow = model_flow(flow, forward=True)
noise = np.random.rand(16, 16, 16, 1)
def feed_dict_fn():
return {images_ph: noise}
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
nets.initialize_actnorms(
sess,
feed_dict_fn=feed_dict_fn,
actnorm_layers=actnorm_layers,
num_steps=50,
)
for actnorm_layer in actnorm_layers:
actnorm_flow = actnorm_layer._forward_outputs[0]
normed_x = sess.run(actnorm_flow[0],
feed_dict={images_ph: noise})
nc = normed_x.shape[-1]
self.assertAllClose(np.var(normed_x.reshape([-1, nc]), axis=0),
[1.0] * nc,
atol=0.1)
self.assertAllClose(np.mean(normed_x.reshape([-1, nc]),
axis=0), [0.0] * nc,
atol=0.1)
def test_simple_affine_coupling_layer(self):
images_np = np.random.rand(8, 32, 32, 16)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.AffineCouplingLayer(
_shift_and_log_scale_fn_template("test"))
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
self.assertEqual(z, None)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
x, logdet = sess.run([x, logdet])
self.assertEqual(x.shape, images_np.shape)
self.assertEqual(logdet.shape, (8, ))
self.forward_inverse(layer, flow)
layer = fl.AffineCouplingLayer(
_shift_and_log_scale_fn_template("train"))
self.try_to_train_identity_layer(layer, flow)
def test_factor_out_layer_conv(self):
images_np = np.random.rand(8, 32, 32, 16)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
layer = fl.FactorOutLayer()
self.forward_inverse(layer, flow)
with tf_framework.arg_scope([fl.FlowLayer.__call__], forward=True):
new_flow = layer(flow)
x, logdet, z = new_flow
self.assertEqual([8, 32, 32, 8], x.shape.as_list())
self.assertEqual([8, 32, 32, 8], z.shape.as_list())
with self.test_session() as sess:
new_flow_np = sess.run(new_flow)
# x
self.assertAllClose(new_flow_np[0], images_np[:, :, :, 8:])
# z
self.assertAllClose(new_flow_np[2], images_np[:, :, :, :8])
# logdet
self.assertAllClose(new_flow_np[1], np.zeros_like(new_flow_np[1]))
def test_actnorm_bias_init_conv_iter(self):
images_ph = tf.placeholder(tf.float32, shape=[8, 32, 32, 3])
flow = fl.InputLayer(images_ph)
layer = fl.ActnormBiasLayer()
new_flow = layer(flow, forward=True)
x, logdet, z = new_flow
init_ops = layer.get_ddi_init_ops(num_init_iterations=100)
with self.test_session() as sess:
# initialize network
sess.run(tf.global_variables_initializer())
for i in range(200):
sess.run(init_ops,
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
# print(sess.run(layer._bias_t))
x_np_values = []
for i in range(20):
x_np, logdet_np = sess.run(
[x, logdet],
feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
x_np_values.append(x_np)
x_np_values = np.array(x_np_values).mean(0)
self.assertEqual(x.shape, x_np_values.shape)
self.assertAllClose(
np.mean(x_np_values.reshape([-1, 3]), axis=0),
[0.0, 0.0, 0.0],
atol=0.05,
)
self.forward_inverse(
layer, flow, feed_dict={images_ph: np.random.rand(8, 32, 32, 3)})
def test_combine_squeeze_and_factor_layers_conv(self):
images_np = np.random.rand(8, 32, 32, 1)
images = tf.to_float(images_np)
flow = fl.InputLayer(images)
# in comments are output shapes
layers = [
fl.SqueezingLayer(), # x=[8, 16, 16, 4]
fl.FactorOutLayer(), # x=[8, 16, 16, 2]
fl.SqueezingLayer(), # x=[8, 8, 8, 8]
fl.FactorOutLayer(), # x=[8, 8, 8, 4] z=[8, 8, 8, 12]
]
chain = fl.ChainLayer(layers=layers)
print()
with tf_framework.arg_scope([fl.FlowLayer.__call__], forward=True):
new_flow = chain(flow)
with self.test_session() as sess:
x, logdet, z = sess.run(new_flow)
self.assertEqual(x.shape, (8, 8, 8, 4))
self.assertEqual(z.shape, (8, 8, 8, 12))
self.forward_inverse(chain, flow)
def model_fn(features, labels, mode, params):
nn_template_fn = nets.OpenAITemplate(
width=args.width
)
layers, actnorm_layers = nets.create_simple_flow(
num_steps=args.num_steps,
num_scales=args.num_scales,
num_bits=args.num_bits,
template_fn=nn_template_fn
)
images = features["images"]
flow = fl.InputLayer(images)
model_flow = fl.ChainLayer(layers)
output_flow = model_flow(flow, forward=True)
y, logdet, z = output_flow
for layer in actnorm_layers:
init_op = layer.get_ddi_init_ops(30)
tf.add_to_collection(ACTNORM_INIT_OPS, init_op)
total_params = 0
trainable_variables = tf.trainable_variables()
for k, v in enumerate(trainable_variables):
num_params = np.prod(v.shape.as_list())
total_params += num_params
print(f"TOTAL PARAMS: {total_params/1e6} [M]")
if mode == tf.estimator.ModeKeys.PREDICT:
raise NotImplementedError()
tfd = tf.contrib.distributions
y_flatten = tf.reshape(y, [batch_size, -1])
z_flatten = tf.reshape(z, [batch_size, -1])
prior_y = tfd.MultivariateNormalDiag(loc=tf.zeros_like(y_flatten),
scale_diag=tf.ones_like(y_flatten))
prior_z = tfd.MultivariateNormalDiag(loc=tf.zeros_like(z_flatten),
scale_diag=tf.ones_like(z_flatten))
log_prob_y = prior_y.log_prob(y_flatten)
log_prob_z = prior_z.log_prob(z_flatten)
loss = log_prob_y + log_prob_z + logdet
# compute loss per pixel, the final loss should be same
# for different input sizes
loss = - tf.reduce_mean(loss) / image_size / image_size
trainable_variables = tf.trainable_variables()
l2_loss = l2_reg * tf.add_n(
[tf.nn.l2_loss(v) for v in trainable_variables])
total_loss = l2_loss + loss
tf.summary.scalar('total_loss', total_loss)
tf.summary.scalar('loss', loss)
tf.summary.scalar('l2_loss', l2_loss)
# Sampling during training and evaluation
prior_y = tfd.MultivariateNormalDiag(loc=tf.zeros_like(y_flatten),
scale_diag=sample_beta * tf.ones_like(y_flatten))
prior_z = tfd.MultivariateNormalDiag(loc=tf.zeros_like(z_flatten),
scale_diag=sample_beta * tf.ones_like(z_flatten))
sample_y_flatten = prior_y.sample()
sample_y = tf.reshape(sample_y_flatten, y.shape.as_list())
sample_z = tf.reshape(prior_z.sample(), z.shape.as_list())
sampled_logdet = prior_y.log_prob(sample_y_flatten)
inverse_flow = sample_y, sampled_logdet, sample_z
sampled_flow = model_flow(inverse_flow, forward=False)
x_flow_sampled, _, _ = sampled_flow
# convert to uint8
quantize_image_layer = layers[0]
x_flow_sampled_uint = quantize_image_layer.to_uint8(x_flow_sampled)
grid_image = tf.contrib.gan.eval.image_grid(
x_flow_sampled_uint,
grid_shape=[4, batch_size // 4],
image_shape=(image_size, image_size),
num_channels=3
)
grid_summary = tf.summary.image(
f'samples{sample_beta}', grid_image, max_outputs=10
)
if mode == tf.estimator.ModeKeys.EVAL:
eval_summary_hook = tf.train.SummarySaverHook(
save_steps=1,
output_dir=args.model_dir + "/eval",
summary_op=grid_summary
)
return tf.estimator.EstimatorSpec(
mode,
loss=total_loss,
evaluation_hooks=[eval_summary_hook]
)
# Create training op.
assert mode == tf.estimator.ModeKeys.TRAIN
train_summary_hook = tf.train.SummarySaverHook(
save_secs=args.save_secs,
output_dir=args.model_dir,
summary_op=grid_summary
)
global_step = tf.train.get_global_step()
learning_rate = tf.train.inverse_time_decay(
args.lr, global_step, args.decay_steps, args.decay_rate,
staircase=True
)
tf.summary.scalar('learning_rate', learning_rate)
optimizer = tf.train.AdagradOptimizer(learning_rate=learning_rate)
if args.clip > 0.0:
gvs = optimizer.compute_gradients(total_loss)
capped_gvs = [
(tf.clip_by_value(grad, -args.clip, args.clip), var) for grad, var in gvs
]
train_op = optimizer.apply_gradients(capped_gvs, global_step=global_step)
else:
train_op = optimizer.minimize(total_loss, global_step=global_step)
return tf.estimator.EstimatorSpec(
mode, loss=total_loss,
train_op=train_op, training_hooks=[train_summary_hook]
)
