def discriminator(imgs, alpha=0.2):
with tf.variable_scope("d", reuse=tf.AUTO_REUSE):
# 14*14*8
layer1 = tf.layers.conv2d(imgs,
filters=8,
kernel_size=3,
strides=(2, 2),
padding="same")
layer1 = tf.maximum(alpha * layer1, layer1)
# 7*7*16
layer2 = tf.layers.conv2d(layer1,
filters=16,
kernel_size=3,
strides=(2, 2),
padding="same")
layer2 = layers.instance_norm(layer2)
layer2 = tf.maximum(alpha * layer2, layer2)
# 3*3*32
layer3 = tf.layers.conv2d(layer2,
filters=32,
kernel_size=3,
strides=(2, 2),
padding="valid")
layer3 = layers.instance_norm(layer3)
layer3 = tf.maximum(alpha * layer3, layer3)
print(layer3)
flatten = tf.reshape(layer3, shape=[-1, 3 * 3 * 32])
logits = tf.layers.dense(flatten, 1)
outputs = tf.sigmoid(logits)
return outputs
def _init_data(self):
self._audio_inputs = self._audio_data.inputs
self._video_inputs = self._video_data.inputs
self._audio_inputs_len = self._audio_data.inputs_length
self._video_inputs_len = self._video_data.inputs_length
# self._labels = self._data.labels
# self._labels_len = self._data.labels_length
if self._hparams.batch_normalisation is True:
self._audio_inputs = tf.layers.batch_normalization(
inputs=self._audio_inputs,
axis=-1,
training=(self._mode == 'train'),
fused=True,
)
self._video_inputs = tf.layers.batch_normalization(
inputs=self._video_inputs,
axis=-1,
training=(self._mode == 'train'),
fused=True,
)
if self._hparams.instance_normalisation is True:
from tensorflow.contrib.layers import instance_norm
self._audio_inputs = instance_norm(inputs=self._audio_inputs, )
self._video_inputs = instance_norm(inputs=self._video_inputs, )
def discriminator(imgs):
with tf.variable_scope("discriminator", reuse=tf.AUTO_REUSE):
conv1 = tf.layers.conv2d(inputs=imgs,
filters=8,
kernel_size=4,
strides=2,
padding="valid")
# in1 = layers.instance_norm(conv1)
leak_relu1 = tf.nn.leaky_relu(conv1)
conv2 = tf.layers.conv2d(inputs=leak_relu1,
filters=16,
kernel_size=4,
strides=2,
padding="valid")
in2 = layers.instance_norm(conv2)
leak_relu2 = tf.nn.leaky_relu(in2)
conv3 = tf.layers.conv2d(inputs=leak_relu2,
filters=32,
kernel_size=4,
strides=2,
padding="valid")
in3 = layers.instance_norm(conv3)
leak_relu3 = tf.nn.leaky_relu(in3)
flatten = tf.layers.flatten(leak_relu3)
logits = tf.layers.dense(flatten, 1)
outputs = tf.nn.sigmoid(logits)
return logits, outputs
def genernator(imgs, is_train=True):
with tf.variable_scope("g", reuse=tf.AUTO_REUSE):
# 28*28*8
layer1 = tf.layers.conv2d(imgs,
filters=8,
kernel_size=3,
strides=(1, 1),
padding="same")
layer1 = layers.instance_norm(layer1)
layer1 = tf.nn.relu(layer1)
# 14*14*16
layer2 = tf.layers.conv2d(layer1,
filters=16,
kernel_size=3,
strides=(2, 2),
padding="same")
layer2 = layers.instance_norm(layer2)
layer2 = tf.nn.relu(layer2)
# 7*7*32
layer3 = tf.layers.conv2d(layer2,
filters=32,
kernel_size=3,
strides=(2, 2),
padding="same")
layer3 = layers.instance_norm(layer3)
layer3 = tf.nn.relu(layer3)
# 4res->7*7*32
res_layer = layer3
for i in range(4):
res_layer = res_block(res_layer, 32, is_train)
# 14*14*16
layer4 = tf.layers.conv2d_transpose(res_layer,
filters=16,
kernel_size=3,
strides=(2, 2),
padding="same")
layer4 = layers.instance_norm(layer4)
layer4 = tf.nn.relu(layer4)
# 28*28*8
layer5 = tf.layers.conv2d_transpose(layer4,
filters=8,
kernel_size=3,
strides=(2, 2),
padding="same")
layer5 = layers.instance_norm(layer5)
layer5 = tf.nn.relu(layer5)
# 28*28*1
logits = tf.layers.conv2d_transpose(layer5,
filters=1,
kernel_size=3,
strides=(1, 1),
padding="same")
outputs = tf.tanh(logits)
return outputs
def conv_instnorm_relu(imgs, filters, kernel_size=3, strides=1):
conv = tf.layers.conv2d(inputs=imgs,
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding="same")
instnorm = layers.instance_norm(conv)
return tf.nn.relu(instnorm)
def transconv_instnorm_relu(imgs, filters, kernel_size=3, strides=2):
transconv = tf.layers.conv2d_transpose(inputs=imgs,
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding="same")
instnorm = layers.instance_norm(transconv)
return tf.nn.relu(instnorm)
def norm(data, is_training, normtype):
if normtype is None:
return data
if normtype.casefold() == 'instance'.casefold():
return layers.instance_norm(data)
if normtype.casefold() == 'batch'.casefold():
return layers.batch_norm(data, is_training=is_training, decay=0.9)
if normtype.casefold() == 'layer'.casefold():
return layers.layer_norm(data)
return data
def __call__(self, x, reuse=tf.AUTO_REUSE, split=False):
with tf.variable_scope(self.name,
reuse=tf.AUTO_REUSE,
regularizer=tcl.l2_regularizer(1e-3)) as vs:
x = tf.reshape(x, [-1, 200, 200, 2])
conv0 = tcl.convolution2d(x,
64, [3, 3], [1, 1],
activation_fn=tf.nn.relu)
conv0 = tcl.convolution2d(conv0,
64, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv0 = tcl.convolution2d(conv0,
128, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv0 = tcl.instance_norm(conv0)
conv1 = tcl.convolution2d(conv0,
128, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv1 = tcl.convolution2d(conv1,
128, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv1 = tcl.convolution2d(conv1,
128, [3, 3], [1, 1],
activation_fn=tf.nn.relu)
conv1 = tcl.instance_norm(conv1)
conv2 = tcl.convolution2d(conv1,
256, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv2 = tcl.convolution2d(conv2,
512, [4, 4], [2, 2],
activation_fn=tf.nn.relu)
conv2 = tcl.convolution2d(conv2,
512, [3, 3], [1, 1],
activation_fn=tf.nn.relu)
conv6 = tcl.flatten(conv2)
fc2 = tcl.fully_connected(conv6,
self.class_num,
activation_fn=tf.identity)
return fc2
def conv_instance_norm_relu(x, W_shape, b_shape):
"""
:param x:
:param W_shape: [H, W, input, output] ([5, 5, 32, 64])
:param b_shape: [output] ([64])
:return:
"""
with tf.name_scope('conv_instance-norm_relu'):
with tf.name_scope('conv2d'):
x = conv2d(x, W_shape, b_shape)
with tf.name_scope('instance-normalization'):
x = layers.instance_norm(x)
with tf.name_scope('relu'):
x = tf.nn.relu(x)
return x
def feature_normalize(x, normalization, name_space, reuse, bn_training, bn_decay):
if normalization == 'W/O' or normalization is None:
y = x
elif normalization == 'batch':
if bn_training is None or bn_decay is None:
raise ValueError('The bn_training and bn_decay must have value'
' when enable Batch Normalization !!!')
y = tf_layers.batch_norm(x,
reuse=reuse,
is_training=bn_training,
decay=bn_decay,
scope=name_space + '/batch_norm')
elif normalization == 'instance':
y = tf_layers.instance_norm(x,
reuse=reuse,
scope=name_space + '/instance_norm')
elif normalization == 'group':
y = tf_layers.group_norm(x,
reuse=reuse,
scope=name_space + '/group_norm')
else:
raise ValueError('Unknown feature normalization method !!!')
return y
def __call__(self, last, reuse=None):
format = self.data_format
# function objects
activation = self.activation
if self.normalization == 'Batch':
normalizer = lambda x: slim.batch_norm(x, 0.999, center=True, scale=True,
is_training=self.training, data_format=format, renorm=False)
elif self.normalization == 'Instance':
normalizer = lambda x: slim.instance_norm(x, center=True, scale=True, data_format=format)
elif self.normalization == 'Group':
normalizer = lambda x: (slim.group_norm(x, x.shape[-3].value // 16, -3, (-2, -1))
if format == 'NCHW' else slim.group_norm(x, x.shape[-1].value // 16, -1, (-3, -2)))
else:
normalizer = None
regularizer = slim.l2_regularizer(self.weight_decay) if self.weight_decay else None
# main model
with tf.variable_scope(self.name, reuse=reuse):
self.training = tf.Variable(False, trainable=False, name='training',
collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.MODEL_VARIABLES])
last = self.def_model(last, activation, normalizer, regularizer)
# trainable/model/save/restore variables
self.tvars = tf.trainable_variables(self.name)
self.mvars = tf.model_variables(self.name)
self.mvars = [var for var in self.mvars if var not in self.tvars]
self.svars = list(set(self.tvars + self.mvars))
self.rvars = self.svars.copy()
# variables that should be affected by weight decay
import re
self.wdvars = [var for var in self.tvars if re.findall(r'weight|kernel', var.name.split('/')[-1])]
# restore moving average of trainable variables
if self.var_ema > 0:
with tf.variable_scope('EMA'):
self.rvars = {**{self.ema.average_name(var): var for var in self.tvars},
**{var.op.name: var for var in self.mvars}}
# return
return last
def create_2D_UNet(x, features_root, n_classes):
net = OrderedDict()
with slim.arg_scope(
[slim.conv2d, slim.conv2d_transpose],
weights_initializer=initializers.variance_scaling_initializer(
factor=2.0, mode='FAN_IN', uniform=False),
activation_fn=leaky_relu):
net['encode/conv1_1'] = instance_norm(
slim.conv2d(x, features_root, [3, 3]))
net['encode/conv1_2'] = instance_norm(
slim.conv2d(net['encode/conv1_1'], features_root, [3, 3]))
net['encode/pool1'] = slim.max_pool2d(net['encode/conv1_2'], [2, 2])
net['encode/conv2_1'] = instance_norm(
slim.conv2d(net['encode/pool1'], features_root * 2, [3, 3]))
net['encode/conv2_2'] = instance_norm(
slim.conv2d(net['encode/conv2_1'], features_root * 2, [3, 3]))
net['encode/pool2'] = slim.max_pool2d(net['encode/conv2_2'], [2, 2])
net['encode/conv3_1'] = instance_norm(
slim.conv2d(net['encode/pool2'], features_root * 4, [3, 3]))
net['encode/conv3_2'] = instance_norm(
slim.conv2d(net['encode/conv3_1'], features_root * 4, [3, 3]))
net['encode/pool3'] = slim.max_pool2d(net['encode/conv3_2'], [2, 2])
net['encode/conv4_1'] = instance_norm(
slim.conv2d(net['encode/pool3'], features_root * 8, [3, 3]))
net['encode/conv4_2'] = instance_norm(
slim.conv2d(net['encode/conv4_1'], features_root * 8, [3, 3]))
net['encode/pool4'] = slim.max_pool2d(net['encode/conv4_2'], [2, 2])
net['encode/conv5_1'] = instance_norm(
slim.conv2d(net['encode/pool4'], features_root * 16, [3, 3]))
net['encode/conv5_2'] = instance_norm(
slim.conv2d(net['encode/conv5_1'], features_root * 16, [3, 3]))
net['decode/up_conv1'] = slim.conv2d_transpose(net['encode/conv5_2'],
features_root * 8,
2,
stride=2,
activation_fn=None,
padding='VALID')
net['decode/concat_c4_u1'] = tf.concat(
[net['encode/conv4_2'], net['decode/up_conv1']], 3)
net['decode/conv1_1'] = instance_norm(
slim.conv2d(net['decode/concat_c4_u1'], features_root * 8, [3, 3]))
net['decode/conv1_2'] = instance_norm(
slim.conv2d(net['decode/conv1_1'], features_root * 8, [3, 3]))
net['decode/up_conv2'] = slim.conv2d_transpose(net['decode/conv1_2'],
features_root * 4,
2,
stride=2,
activation_fn=None,
padding='VALID')
net['decode/concat_c3_u2'] = tf.concat(
[net['encode/conv3_2'], net['decode/up_conv2']], 3)
net['decode/conv2_1'] = instance_norm(
slim.conv2d(net['decode/concat_c3_u2'], features_root * 4, [3, 3]))
net['decode/conv2_2'] = instance_norm(
slim.conv2d(net['decode/conv2_1'], features_root * 4, [3, 3]))
net['decode/up_conv3'] = slim.conv2d_transpose(net['decode/conv2_2'],
features_root * 2,
2,
stride=2,
activation_fn=None,
padding='VALID')
net['decode/concat_c2_u3'] = tf.concat(
[net['encode/conv2_2'], net['decode/up_conv3']], 3)
net['decode/conv3_1'] = instance_norm(
slim.conv2d(net['decode/concat_c2_u3'], features_root * 2, [3, 3]))
net['decode/conv3_2'] = instance_norm(
slim.conv2d(net['decode/conv3_1'], features_root * 2, [3, 3]))
net['decode/up_conv4'] = slim.conv2d_transpose(net['decode/conv3_2'],
features_root,
2,
stride=2,
activation_fn=None,
padding='VALID')
net['decode/concat_c1_u4'] = tf.concat(
[net['encode/conv1_2'], net['decode/up_conv4']], 3)
net['decode/conv4_1'] = instance_norm(
slim.conv2d(net['decode/concat_c1_u4'], features_root, [3, 3]))
net['decode/conv4_2'] = instance_norm(
slim.conv2d(net['decode/conv4_1'], features_root, [3, 3]))
net['out_map'] = instance_norm(
slim.conv2d(net['decode/conv4_2'],
n_classes, [1, 1],
activation_fn=None))
return net
def forward(x, is_training):
return instance_norm(x,
center=True,
scale=True,
epsilon=1e-06,
trainable=is_training)
def _instance_norm(input):
with tf.variable_scope('instance_norm'):
return li.instance_norm(input, center=True, scale=True, epsilon=1e-5)
def feature_extractor_c(self, reuse=False):
training = tf.cond(self.train, lambda: True, lambda: False)
X = layers.instance_norm(self.X)
with tf.variable_scope('feature_extractor_c', reuse=reuse):
h_conv1 = layers.conv2d(self.X,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.conv2d(h_conv1,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.conv2d(h_conv1,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.max_pool2d(h_conv1, 2, 2, padding='SAME')
h_conv1 = noise(tf.layers.dropout(h_conv1,
rate=0.5,
training=training),
phase=training)
h_conv2 = layers.conv2d(h_conv1,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.max_pool2d(h_conv2, 2, 2, padding='SAME')
h_conv2 = noise(tf.layers.dropout(h_conv2,
rate=0.5,
training=training),
phase=training)
h_conv3 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = layers.conv2d(h_conv3,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = layers.conv2d(h_conv3,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = tf.reduce_mean(h_conv3, axis=[1, 2])
self.features_c = h_conv3
feature_shape = self.features_c.get_shape()
self.feature_c_dim = feature_shape[1].value
self.features_c_src = tf.slice(self.features_c, [0, 0],
[self.batch_size, -1])
self.features_c_for_prediction = tf.cond(
self.train, lambda: tf.slice(self.features_c, [0, 0],
[self.batch_size, -1]),
lambda: self.features_c)
def sn_spade_normalize(tensor,
condition,
second_condition=None,
scope="spade",
is_training=False):
"""A spectral normalized version of SPADE.
Performs SPADE normalization (Park et al.) on a tensor based on condition. If
second_condition is defined, concatenate second_condition to condition.
These inputs are separated because they encode conditioning of different
things.
Args:
tensor: [B, H, W, D] a tensor to apply SPADE normalization to
condition: [B, H', W', D'] A tensor used to predict SPADE's normalization
parameters.
second_condition: [B, H'', W'', D''] A tensor used to encode another kind of
conditioning. second_condition is provided in case its dimensions do not
natively match condition's dimension.
scope: (str) The scope of the SPADE convolutions.
is_training: (bool) used to control the spectral normalization update
schedule. When true apply an update, else freeze updates.
Returns:
A SPADE normalized tensor of shape [B, H, W, D].
"""
# resize condition to match input spatial
n_tensor = layers.instance_norm(tensor,
center=False,
scale=False,
trainable=False,
epsilon=1e-4)
unused_b, h_tensor, w_tensor, feature_dim = n_tensor.shape.as_list()
with tf.compat.v1.variable_scope(scope, reuse=tf.AUTO_REUSE):
resize_condition = diff_resize_area(condition, [h_tensor, w_tensor])
if second_condition is not None:
second_condition = diff_resize_area(second_condition,
[h_tensor, w_tensor])
resize_condition = tf.concat([resize_condition, second_condition],
axis=-1)
resize_condition = utils.pad_panorama_for_convolutions(
resize_condition, 3, "symmetric")
condition_net = tf.nn.relu(
sn_ops.snconv2d(resize_condition,
32,
3,
1,
is_training=is_training,
name="intermediate_spade"))
condition_net = utils.pad_panorama_for_convolutions(
condition_net, 3, "symmetric")
gamma_act = sn_ops.snconv2d(condition_net,
feature_dim,
3,
1,
is_training=is_training,
name="g_spade")
mu_act = sn_ops.snconv2d(condition_net,
feature_dim,
3,
1,
is_training=is_training,
name="b_spade")
return n_tensor * (1 + gamma_act) - mu_act
def __call__(self, last, reuse=None):
format = self.data_format
kernel1 = [1, 8]
stride1 = [1, 2]
kernel2 = [1, 3]
stride2 = [1, 2]
# function objects
activation = self.activation
if self.normalization == 'Batch':
normalizer = lambda x: slim.batch_norm(x,
0.999,
center=True,
scale=True,
is_training=self.training,
data_format=format,
renorm=False)
elif self.normalization == 'Instance':
normalizer = lambda x: slim.instance_norm(
x, center=True, scale=True, data_format=format)
elif self.normalization == 'Group':
normalizer = lambda x: (slim.group_norm(
x,
x.shape.as_list()[-3] // 16, -3,
(-2, -1)) if format == 'NCHW' else slim.group_norm(
x,
x.shape.as_list()[-1] // 16, -1, (-3, -2)))
else:
normalizer = None
regularizer = slim.l2_regularizer(
self.weight_decay) if self.weight_decay else None
skip_connection = lambda x, y: x + y
# skip_connection = lambda x, y: tf.concat([x, y], -3 if format == 'NCHW' else -1)
# model scope
with tf.variable_scope(self.name, reuse=reuse):
# states
self.training = tf.Variable(False,
trainable=False,
name='training',
collections=[
tf.GraphKeys.GLOBAL_VARIABLES,
tf.GraphKeys.MODEL_VARIABLES
])
skips = []
# encoder
with tf.variable_scope('InBlock'):
last = self.EBlock(last, 16, 0, [1, 8], [1, 1], format, None,
None, regularizer)
with tf.variable_scope('EBlock_0'):
skips.append(last)
last = self.EBlock(last, 32, 0, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_1'):
skips.append(last)
last = self.EBlock(last, 48, 0, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_2'):
skips.append(last)
last = self.EBlock(last, 64, 1, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_3'):
skips.append(last)
last = self.EBlock(last, 96, 1, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_4'):
skips.append(last)
last = self.EBlock(last, 128, 2, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_5'):
skips.append(last)
last = self.EBlock(last, 160, 2, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_6'):
skips.append(last)
last = self.EBlock(last, 192, 2, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_7'):
skips.append(last)
last = self.EBlock(last, 224, 3, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_8'):
skips.append(last)
last = self.EBlock(last, 256, 3, kernel1, stride1, format,
activation, normalizer, regularizer)
# decoder
with tf.variable_scope('DBlock_8'):
last = self.DBlock(last, 224, 3, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_7'):
last = self.DBlock(last, 192, 2, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_6'):
last = self.DBlock(last, 160, 2, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_5'):
last = self.DBlock(last, 128, 2, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_4'):
last = self.DBlock(last, 96, 1, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_3'):
last = self.DBlock(last, 64, 1, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_2'):
last = self.DBlock(last, 48, 0, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_1'):
last = self.DBlock(last, 32, 0, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('DBlock_0'):
last = self.DBlock(last, 16, 0, kernel2, stride2, format,
activation, normalizer, regularizer)
last = skip_connection(last, skips.pop())
with tf.variable_scope('OutBlock'):
last = self.EBlock(last, self.in_channels, 0, [1, 8], [1, 1],
format, activation, normalizer, regularizer)
# trainable/model/save/restore variables
self.tvars = tf.trainable_variables(self.name)
self.mvars = tf.model_variables(self.name)
self.mvars = [i for i in self.mvars if i not in self.tvars]
self.svars = list(set(self.tvars + self.mvars))
self.rvars = self.svars.copy()
# restore moving average of trainable variables
if self.var_ema > 0:
with tf.variable_scope('EMA'):
self.rvars = {
**{self.ema.average_name(var): var
for var in self.tvars},
**{var.op.name: var
for var in self.mvars}
}
return last
def instance_norm(self, *args, **kwargs):
return layers.instance_norm(args[0], kwargs["name"])
def instance_norm(input, name='instance_norm'):
return layers.instance_norm(input, scope=name)
def decoder(encoded,
scales,
styles,
texture_only=False,
style_size=8,
image_size=(112, 112),
keep_prob=1.0,
phase_train=True,
weight_decay=0.0,
reuse=None,
scope='Decoder'):
with tf.variable_scope(scope, reuse=reuse):
with slim.arg_scope(
[slim.conv2d, slim.conv2d_transpose, slim.fully_connected],
activation_fn=tf.nn.relu,
# weights_initializer=tf.contrib.layers.xavier_initializer(),
weights_initializer=tf.contrib.layers.
variance_scaling_initializer(),
weights_regularizer=slim.l2_regularizer(weight_decay)):
with slim.arg_scope([slim.dropout, slim.batch_norm],
is_training=phase_train):
with slim.arg_scope([slim.fully_connected],
normalizer_fn=layer_norm,
normalizer_params=None):
print('{} input shape:'.format(scope),
[dim.value for dim in encoded.shape])
batch_size = tf.shape(encoded)[0]
h, w = tuple(image_size)
k = 64
with tf.variable_scope('StyleController'):
if styles is None:
styles = tf.random_normal((batch_size, style_size))
net = tf.identity(styles, name='input_style')
net = slim.fully_connected(net, 128, scope='fc2')
print('module fc2 shape:',
[dim.value for dim in net.shape])
net = slim.fully_connected(net, 128, scope='fc3')
print('module fc3 shape:',
[dim.value for dim in net.shape])
gamma = slim.fully_connected(net,
4 * k,
activation_fn=None,
normalizer_fn=None,
scope='fc4')
gamma = tf.reshape(gamma, [-1, 1, 1, 4 * k],
name='gamma')
print('gamma shape:',
[dim.value for dim in gamma.shape])
beta = slim.fully_connected(net,
4 * k,
activation_fn=None,
normalizer_fn=None,
scope='fc5')
beta = tf.reshape(beta, [-1, 1, 1, 4 * k], name='beta')
print('beta shape:', [dim.value for dim in beta.shape])
with tf.variable_scope('Decoder'):
print('-- Decoder')
net = encoded
adain = lambda x: gamma * instance_norm(
x, center=False, scale=False) + beta
with slim.arg_scope(
[slim.conv2d_transpose, slim.conv2d],
normalizer_fn=adain,
normalizer_params=None):
for i in range(3):
net_ = conv(net,
4 * k,
3,
scope='res{}_0'.format(i))
net += conv(net_,
4 * k,
3,
activation_fn=None,
biases_initializer=None,
scope='res{}_1'.format(i))
print('module res{} shape:'.format(i),
[dim.value for dim in net.shape])
with slim.arg_scope([
slim.conv2d, slim.conv2d_transpose,
slim.fully_connected
],
normalizer_fn=layer_norm,
normalizer_params=None):
net = upscale2d(net, 2)
net = conv(net, 2 * k, 5, pad=2, scope='deconv1_1')
print('module deconv1 shape:',
[dim.value for dim in net.shape])
net = upscale2d(net, 2)
net = conv(net, k, 5, pad=2, scope='deconv2_1')
net = conv(
net,
3,
7,
pad=3,
activation_fn=None,
normalizer_fn=None,
weights_initializer=tf.constant_initializer(0.0),
scope='conv_image')
images_rendered = tf.nn.tanh(net,
name='images_rendered')
print('images_rendered shape:',
[dim.value for dim in images_rendered.shape])
if texture_only:
return images_rendered
with tf.variable_scope('WarpController'):
print('-- WarpController')
net = encoded
warp_input = tf.identity(images_rendered,
name='warp_input')
net = slim.flatten(net)
net = slim.fully_connected(net, 128, scope='fc1')
print('module fc1 shape:',
[dim.value for dim in net.shape])
num_ldmark = 16
# Predict the control points
ldmark_mean = (
np.random.normal(0, 50, (num_ldmark, 2)) +
np.array([[0.5 * h, 0.5 * w]])).flatten()
ldmark_mean = tf.Variable(ldmark_mean.astype(
np.float32),
name='ldmark_mean')
print('ldmark_mean shape:',
[dim.value for dim in ldmark_mean.shape])
ldmark_pred = slim.fully_connected(
net,
num_ldmark * 2,
weights_initializer=tf.
truncated_normal_initializer(stddev=1.0),
normalizer_fn=None,
activation_fn=None,
biases_initializer=None,
scope='fc_ldmark')
ldmark_pred = ldmark_pred + ldmark_mean
print('ldmark_pred shape:',
[dim.value for dim in ldmark_pred.shape])
ldmark_pred = tf.identity(ldmark_pred,
name='ldmark_pred')
# Predict the displacements
ldmark_diff = slim.fully_connected(net,
num_ldmark * 2,
normalizer_fn=None,
activation_fn=None,
scope='fc_diff')
print('ldmark_diff shape:',
[dim.value for dim in ldmark_diff.shape])
ldmark_diff = tf.identity(ldmark_diff,
name='ldmark_diff')
ldmark_diff = tf.identity(tf.reshape(scales, [-1, 1]) *
ldmark_diff,
name='ldmark_diff_scaled')
src_pts = tf.reshape(ldmark_pred, [-1, num_ldmark, 2])
dst_pts = tf.reshape(ldmark_pred + ldmark_diff,
[-1, num_ldmark, 2])
diff_norm = tf.reduce_mean(
tf.norm(src_pts - dst_pts, axis=[1, 2]))
# tf.summary.scalar('diff_norm', diff_norm)
# tf.summary.scalar('mark', ldmark_pred[0,0])
images_transformed, dense_flow = sparse_image_warp(
warp_input,
src_pts,
dst_pts,
regularization_weight=1e-6,
num_boundary_points=0)
dense_flow = tf.identity(dense_flow, name='dense_flow')
return images_transformed, images_rendered, ldmark_pred, ldmark_diff
def feature_extractor_d(self, reuse=False):
training = tf.cond(self.train, lambda: True, lambda: False)
X = layers.instance_norm(self.X)
with tf.variable_scope('feature_extractor_d', reuse=reuse):
h_conv1 = layers.conv2d(self.X,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.conv2d(h_conv1,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.conv2d(h_conv1,
self.ef_dim * 3,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv1 = layers.batch_norm(h_conv1, activation_fn=leaky_relu)
h_conv1 = layers.max_pool2d(h_conv1, 2, 2, padding='SAME')
h_conv1 = noise(tf.layers.dropout(h_conv1,
rate=0.5,
training=training),
phase=training)
h_conv2 = layers.conv2d(h_conv1,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv2 = layers.batch_norm(h_conv2, activation_fn=leaky_relu)
h_conv2 = layers.max_pool2d(h_conv2, 2, 2, padding='SAME')
h_conv2 = noise(tf.layers.dropout(h_conv2,
rate=0.5,
training=training),
phase=training)
h_conv3 = layers.conv2d(h_conv2,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = layers.conv2d(h_conv3,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = layers.conv2d(h_conv3,
self.ef_dim * 6,
3,
stride=1,
padding='SAME',
activation_fn=None,
weights_initializer=self.initializer)
h_conv3 = layers.batch_norm(h_conv3, activation_fn=leaky_relu)
h_conv3 = tf.reduce_mean(h_conv3, axis=[1, 2])
self.features_d = h_conv3
def instance_norm(x, scope='instance_normalization'):
return layers.instance_norm(x, epsilon=1e-05, center=True, scale=True, scope=scope)
def instance_norm_tf(x, training=False, reuse=False, scope=None):
return instance_norm(x, epsilon=1e-5, reuse=False, scope=scope)
def __call__(self, last, reuse=None):
# parameters
format = self.data_format
kernel1 = [3, 3]
stride1 = [2, 2]
# function objects
activation = self.activation
if self.normalization == 'Batch':
normalizer = lambda x: slim.batch_norm(x,
0.999,
center=True,
scale=True,
is_training=self.training,
data_format=format,
renorm=False)
elif self.normalization == 'Instance':
normalizer = lambda x: slim.instance_norm(
x, center=True, scale=True, data_format=format)
elif self.normalization == 'Group':
normalizer = lambda x: (slim.group_norm(
x,
x.shape.as_list()[-3] // 16, -3,
(-2, -1)) if format == 'NCHW' else slim.group_norm(
x,
x.shape.as_list()[-1] // 16, -1, (-3, -2)))
else:
normalizer = None
regularizer = slim.l2_regularizer(
self.weight_decay) if self.weight_decay else None
# model scope
with tf.variable_scope(self.name, reuse=reuse):
# states
self.training = tf.Variable(False,
trainable=False,
name='training',
collections=[
tf.GraphKeys.GLOBAL_VARIABLES,
tf.GraphKeys.MODEL_VARIABLES
])
# encoder
with tf.variable_scope('InBlock'):
last = self.InBlock(last, 32, [7, 7], [1, 1], format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_1'):
last = self.EBlock(last, 1, 48, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_2'):
last = self.EBlock(last, 1, 64, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_3'):
last = self.EBlock(last, 1, 96, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_4'):
last = self.EBlock(last, 1, 128, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_5'):
last = self.EBlock(last, 1, 192, kernel1, stride1, format,
activation, normalizer, regularizer)
with tf.variable_scope('EBlock_6'):
last = self.EBlock(last, 1, 256, kernel1, [1, 1], format,
activation, normalizer, regularizer)
with tf.variable_scope('PatchCritic'):
last_channels = last.shape.as_list()[-3]
critic_logit = self.ResBlock(last,
last_channels, [3, 3], [1, 1],
format=format,
activation=activation,
normalizer=normalizer,
regularizer=regularizer)
critic_logit = slim.conv2d(critic_logit,
1, [7, 7], [1, 1],
'SAME',
format,
1,
None,
None,
weights_regularizer=regularizer)
with tf.variable_scope('GlobalAveragePooling'):
last = tf.reduce_mean(
last, [-2, -1] if format == 'NCHW' else [-3, -2])
with tf.variable_scope('Domain'):
last_channels = last.shape.as_list()[-1]
domain_logit = slim.fully_connected(
last,
last_channels,
activation,
None,
weights_regularizer=regularizer)
domain_logit = slim.fully_connected(
domain_logit,
self.num_domains,
None,
None,
weights_regularizer=regularizer)
# trainable/model/save/restore variables
self.tvars = tf.trainable_variables(self.name)
self.mvars = tf.model_variables(self.name)
self.mvars = [i for i in self.mvars if i not in self.tvars]
self.svars = list(set(self.tvars + self.mvars))
self.rvars = self.svars.copy()
return critic_logit, domain_logit
