def __init__(self):
nn_template_fn = nets.OpenAITemplate(width=c.WIDTH_RESNET)
layers, self.actnorm_layers = nets.create_simple_flow(
num_steps=c.N_FLOW_STEPS,
num_scales=c.N_FLOW_SCALES,
template_fn=nn_template_fn)
self.model_flow = fl.ChainLayer(layers)
self.quantize_image_layer = layers[0]
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 create_simple_flow(
num_steps: int = 1,
num_scales: int = 3,
num_bits: int = 5,
template_fn: Any = ResentTemplate()
) -> Tuple[List[fl.FlowLayer], List[fl.ActnormLayer]]:
"""Create Glow model. This implementation may slightly differ from the
official one. For example the last layer here is the fl.FactorOutLayer
Args:
num_steps: number of steps per single scale, a K parameter from the paper
num_scales: number of scales, a L parameter from the paper. Each scale
reduces the tensor spatial dimension by 2.
num_bits: input image quantization
template_fn: a template function used in AffineCoupling layer
Returns:
layers: a list of layers which define normalizing flow
actnorms: a list of actnorm layers which can be initialized using data
dependent initialization. See: initialize_actnorms() function.
"""
layers = [fl.QuantizeImage(num_bits=num_bits)]
actnorm_layers = []
for scale in range(num_scales):
scale_name = f"Scale{scale+1}"
scale_steps = []
for s in range(num_steps):
name = f"Step{s+1}"
step_layer, actnorm_layer = step_flow(
name=name,
shift_and_log_scale_fn=template_fn.create_template_fn(name))
scale_steps.append(step_layer)
actnorm_layers.append(actnorm_layer)
layers += [
fl.SqueezingLayer(name=scale_name),
fl.ChainLayer(scale_steps, name=scale_name),
fl.FactorOutLayer(name=scale_name),
]
return layers, actnorm_layers
def step_flow(
name: str, shift_and_log_scale_fn: Callable[[tf.Tensor], tf.Tensor]
) -> Tuple[fl.ChainLayer, fl.ActnormLayer]:
"""Create single step of the Glow model:
1. actnorm
2. invertible conv
3. affine coupling layer
Returns:
step_layer: a flow layer which perform 1-3 operations
actnorm: a reference of actnorm layer from step 1. This reference can be
used to initialize this layer using data dependent initialization
"""
actnorm = fl.ActnormLayer()
layers = [
actnorm,
fl.InvertibleConv1x1Layer(),
fl.AffineCouplingLayer(shift_and_log_scale_fn=shift_and_log_scale_fn),
]
return fl.ChainLayer(layers, name=name), actnorm
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_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]
)