def block(x, labels, out_channels, num_classes, is_training, CGN, CGN_groups,
name):
with tf.variable_scope(name):
if CGN:
norm0 = ops.ConditionalGroupNorm(num_classes,
CGN_groups,
name='cgn_0')
norm1 = ops.ConditionalGroupNorm(num_classes,
CGN_groups,
name='cgn_1')
else:
norm0 = ops.ConditionalBatchNorm(num_classes, name='cbn_0')
norm1 = ops.ConditionalBatchNorm(num_classes, name='cbn_1')
x_0 = x
x = tf.nn.relu(norm0(x, labels, is_training))
x = usample(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv1')
x = tf.nn.relu(norm1(x, labels, is_training))
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv2')
x_0 = usample(x_0)
x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='snconv3')
return x_0 + x
def _d_residual_block(self,
x,
out_ch,
idx,
is_training,
resize=True,
is_head=False):
update_collection = self._get_update_collection(is_training)
with tf.variable_scope("d_resblock_" + str(idx), reuse=tf.AUTO_REUSE):
h = x
if not is_head:
h = tf.nn.relu(h)
h = snconv2d(h,
out_ch,
name='d_resblock_conv_1',
update_collection=update_collection)
h = tf.nn.relu(h)
h = snconv2d(h,
out_ch,
name='d_resblock_conv_2',
update_collection=update_collection)
if resize:
h = slim.avg_pool2d(h, [2, 2])
# Short cut
s = x
if resize:
s = slim.avg_pool2d(s, [2, 2])
s = snconv2d(s,
out_ch,
k_h=1,
k_w=1,
name='d_resblock_conv_sc',
update_collection=update_collection)
return h + s
def block(x, labels, out_channels, num_classes, name, training=True):
"""Builds the residual blocks used in the generator.
Args:
x: The 4D input tensor.
labels: The labels of the class we seek to sample from.
out_channels: Integer number of features in the output layer.
num_classes: Integer number of classes in the labels.
name: The variable scope name for the block.
training: Whether this block is for training or not.
Returns:
A `Tensor` representing the output of the operation.
"""
with tf.compat.v1.variable_scope(name):
labels_onehot = tf.one_hot(labels, num_classes)
x_0 = x
x = tf.nn.relu(
tfgan.tpu.batch_norm(x, training, labels_onehot, name='cbn_0'))
x = usample(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv1')
x = tf.nn.relu(
tfgan.tpu.batch_norm(x, training, labels_onehot, name='cbn_1'))
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv2')
x_0 = usample(x_0)
x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, training, 'snconv3')
return x_0 + x
def _g_residual_block(self, x, y, n_ch, idx, is_training, resize=True):
update_collection = self._get_update_collection(is_training)
with tf.variable_scope("g_resblock_" + str(idx), reuse=tf.AUTO_REUSE):
h = self._cbn(x, y, is_training, scope='g_resblock_cbn_1')
h = tf.nn.relu(h)
if resize:
h = upscale(h, 2)
h = snconv2d(h,
n_ch,
name='g_resblock_conv_1',
update_collection=update_collection)
h = self._cbn(h, y, is_training, scope='g_resblock_cbn_2')
h = tf.nn.relu(h)
h = snconv2d(h,
n_ch,
name='g_resblock_conv_2',
update_collection=update_collection)
if resize:
sc = upscale(x, 2)
else:
sc = x
sc = snconv2d(sc,
n_ch,
k_h=1,
k_w=1,
name='g_resblock_conv_sc',
update_collection=update_collection)
return h + sc
def block(x, out_channels, name, downsample=True, act=tf.nn.relu):
"""Builds the residual blocks used in the discriminator.
Args:
x: The 4D input vector.
out_channels: Number of features in the output layer.
name: The variable scope name for the block.
downsample: If True, downsample the spatial size the input tensor by
a factor of 2 on each side. If False, the spatial size of the
input tensor is unchanged.
act: The activation function used in the block.
Returns:
A `Tensor` representing the output of the operation.
"""
with tf.compat.v1.variable_scope(name):
input_channels = x.shape.as_list()[-1]
x_0 = x
x = act(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv1')
x = act(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv2')
if downsample:
x = dsample(x)
if downsample or input_channels != out_channels:
x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='sn_conv3')
if downsample:
x_0 = dsample(x_0)
return x_0 + x
def block_without_condition(x, out_channels, is_training, name):
with tf.variable_scope(name):
bn0 = ops.batch_norm(name='bn_0')
bn1 = ops.batch_norm(name='bn_1')
# bn = ops.batch_norm()
x = bn0(x, is_training)
x = tf.nn.relu(x)
x = usample(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv1')
x = bn1(x, is_training)
x = tf.nn.relu(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv2')
return x
def optimized_block(x,
out_channels,
name,
update_collection=None,
act=tf.nn.relu):
"""Builds the simplified residual blocks for downsampling.
Compared with block, optimized_block always downsamples the spatial resolution
of the input vector by a factor of 4.
Args:
x: The 4D input vector.
out_channels: Number of features in the output layer.
name: The variable scope name for the block.
update_collection: The update collections used in the
spectral_normed_weight.
act: The activation function used in the block.
Returns:
A `Tensor` representing the output of the operation.
"""
with tf.variable_scope(name):
x_0 = x
x = ops.snconv2d(x,
out_channels,
3,
3,
1,
1,
update_collection=update_collection,
name='sn_conv1')
x = act(x)
x = ops.snconv2d(x,
out_channels,
3,
3,
1,
1,
update_collection=update_collection,
name='sn_conv2')
x = dsample(x)
x_0 = dsample(x_0)
x_0 = ops.snconv2d(x_0,
out_channels,
1,
1,
1,
1,
update_collection=update_collection,
name='sn_conv3')
return x + x_0
def block(x, labels, out_channels, num_classes, is_training, name):
with tf.variable_scope(name):
bn0 = ops.ConditionalBatchNorm(num_classes, name='cbn_0')
bn1 = ops.ConditionalBatchNorm(num_classes, name='cbn_1')
x_0 = x
x = tf.nn.relu(bn0(x, labels, is_training))
x = usample(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv1')
x = tf.nn.relu(bn1(x, labels, is_training))
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='snconv2')
x_0 = usample(x_0)
x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='snconv3')
return x_0 + x
def generator(z,trainable=True, reuse=tf.AUTO_REUSE):
#z = tf.reshape(z,(-1,1,1,n_dim))
with tf.variable_scope("generator", reuse=reuse):
ch = 1024
z = ops.linear(z,ch*4*4,scope='g_h0')
z = tf.reshape(z,(-1,4,4,ch))#4*4*1024
print(z)
z = risidual_up_block(z,ch,trainable,scope='deconv0')#8*8*1024
print(z)
z = risidual_up_block(z,ch//2,trainable,scope='deconv1')#16*16*512
print(z)
z = risidual_up_block(z,ch//4,trainable,scope='deconv2')#32*32*256
print(z)
z = risidual_up_block(z,ch//8,trainable,scope='deconv3')#64*64*128
z = attention(z,z.shape[-1])
print(z)
z = risidual_up_block(z,ch//16,trainable,scope='deconv4')#128*128*64
print(z)
z = risidual_up_block(z,ch//32,trainable,scope='deconv5')#256*256*64
print(z)
z = tf.layers.batch_normalization(z,training=trainable)
z = tf.nn.relu(z)
z = ops.snconv2d(z,channel,3,3,1,1,name='last_layer')
z = tf.nn.tanh(z)
print(z)
return z
def generator_without_condition(zs, gf_dim, is_training=True):
"""Builds the generator graph propagating from z to x.
Args:
zs: The list of noise tensors.
gf_dim: The gf dimension.
scope: Optional scope for `variable_op_scope`.
is_training: is training
Returns:
outputs: The output layer of the generator.
"""
with tf.variable_scope('model', reuse=tf.AUTO_REUSE):
# project `z` and reshape
act0 = ops.snlinear(zs, gf_dim * 16 * 4 * 4, name='g_snh0')
act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])
act1 = block_without_condition(act0, gf_dim * 16, is_training, 'g_block1') # 8 * 8
act2 = block_without_condition(act1, gf_dim * 8, is_training, 'g_block2') # 16 * 16
act3 = block_without_condition(act2, gf_dim * 4, is_training, 'g_block3') # 32 * 32
# act3 = non_local.sn_non_local_block_sim(act3, None, name='g_non_local')
act4 = block_without_condition(act3, gf_dim * 2, is_training, 'g_block4') # 64 * 64
act4, attn = non_local.sn_non_local_block_sim(act4, None, name='g_non_local')
act5 = block_without_condition(act4, gf_dim, is_training, 'g_block5') # 128 * 128
bn = ops.batch_norm(name='g_bn')
act5 = tf.nn.relu(bn(act5, is_training))
act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, name='g_snconv_last')
out = tf.nn.tanh(act6)
print('Generator without Condition with SA')
return out, attn
def discriminator(x,trainable=True,reuse=tf.AUTO_REUSE):
with tf.variable_scope("discriminator", reuse=reuse):
ch = 32
x = ops.snconv2d(x,16,7,7,1,1,name='init_block')
#x = tf.nn.leaky_relu(x)
x = risidual_down_block(x,ch,'conv0')#128*128*64
print(x)
x = risidual_down_block(x,ch*2,'conv1')#64*64*128
print(x)
x = risidual_down_block(x,ch*4,'conv2')#32*32*256
x = attention(x,x.shape[-1])
print(x)
x = risidual_down_block(x,ch*8,'conv3') #16*16*512
print(x)
x = risidual_down_block(x,ch*16,'conv4') #8*8*1024
print(x)
x = risidual_down_block(x,ch*32,'conv5') #4*4*1024
print(x)
x = tf.reduce_sum(x, [1, 2])
print(x)
x = ops.snlinear(x,1,name='d_sn_linear')
print(x)
return x
def forward(self, z, coord, is_training):
valid_sizes = {4, 8, 16, 32, 64, 128, 256}
assert (self.micro_patch_size[0] in valid_sizes and self.micro_patch_size[1] in valid_sizes), \
"I haven't test your micro patch size: {}".format(self.micro_patch_size)
update_collection = self._get_update_collection(is_training)
print(" [Build] Generator ; is_training: {}".format(is_training))
with tf.variable_scope("G_generator", reuse=tf.AUTO_REUSE):
init_sp = 2
init_ngf_mult = 16
cond = tf.concat([z, coord], axis=1)
h = snlinear(cond,
self.ngf_base * init_ngf_mult * init_sp * init_sp,
'g_z_fc',
update_collection=update_collection)
h = tf.reshape(
h, [-1, init_sp, init_sp, self.ngf_base * init_ngf_mult])
# Stacking residual blocks
num_resize_layers = int(
math.log(min(self.micro_patch_size), 2) - 1)
num_total_layers = num_resize_layers + self.num_extra_layers
basic_layers = [8, 4, 2]
if num_total_layers >= len(basic_layers):
num_replicate_layers = num_total_layers - len(basic_layers)
ngf_mult_list = basic_layers + [
1,
] * num_replicate_layers
else:
ngf_mult_list = basic_layers[:num_total_layers]
print("\t ngf_mult_list = {}".format(ngf_mult_list))
for idx, ngf_mult in enumerate(ngf_mult_list):
n_ch = self.ngf_base * ngf_mult
# Standard layers first
if idx < num_resize_layers:
resize, is_extra = True, False
# Extra layers do not resize spatial size
else:
resize, is_extra = False, True
h = self._g_residual_block(h,
cond,
n_ch,
idx=idx,
is_training=is_training,
resize=resize)
print(
"\t GResBlock: id={}, out_shape={}, resize={}, is_extra={}"
.format(idx, h.shape.as_list(), resize, is_extra))
h = batch_norm(name="g_last_bn")(h, is_training=is_training)
h = tf.nn.relu(h)
h = snconv2d(h,
self.c_dim,
name='g_last_conv_2',
update_collection=update_collection)
return tf.nn.tanh(h)
def risidual_down_block(x,ch,scope):
with tf.variable_scope(scope):
short_cut = x
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch//4,1,1,1,1,name='sn_conv0')
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch, 3, 3, 1, 1,name='sn_conv1')
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch,3,3,1,1,name='sn_conv2')
x = tf.nn.relu(x)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
x = ops.snconv2d(x,ch,1,1,1,1,name='sn_conv3')
short_cut = tf.nn.avg_pool(short_cut, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
short_cut2 = ops.snconv2d(short_cut,short_cut.shape[-1],1,1,1,1,name='sc_sn_conv')
short_cut = Concatenation([short_cut,short_cut2])
return x+short_cut
def attention(x, ch, scope='attention', reuse=tf.AUTO_REUSE):
with tf.variable_scope(scope, reuse=reuse):
f = ops.snconv2d(x,ch//8,1,1,1,1,name='f_conv')
g = ops.snconv2d(x,ch//8,1,1,1,1,name='g_conv')
h = ops.snconv2d(x,ch,1,1,1,1,name='h_conv')
f = tf.reshape(f,(batch_size,-1,f.shape[-1]))
g = tf.reshape(g,(batch_size,-1,g.shape[-1]))
h = tf.reshape(h,(batch_size,-1,h.shape[-1]))
s = tf.matmul(g,f,transpose_b=True)
beta = tf.nn.softmax(s)
o = tf.matmul(beta,h)
gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0))
o = tf.reshape(o,shape=x.shape)
o = ops.snconv2d(o, ch,1,1,1,1,name='attn_conv')
x = gamma*o+x
return x
def risidual_up_block(x,ch,trainable,scope=''):
with tf.variable_scope(scope):
short_cut = x
x = tf.layers.batch_normalization(x,training=trainable)
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch//4,1,1,1,1,name='sn_upconv0')
x = tf.layers.batch_normalization(x,training=trainable)
x = tf.nn.relu(x)
x = upsample(x)
x = ops.snconv2d(x,ch,3,3,1,1,name='sn_upconv1')
x = tf.layers.batch_normalization(x,training=trainable)
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch,3,3,1,1,name='sn_upconv2')
x = tf.layers.batch_normalization(x,training=trainable)
x = tf.nn.relu(x)
x = ops.snconv2d(x,ch,1,1,1,1,name='sn_upconv3')
short_cut = upsample(short_cut)
short_cut = ops.snconv2d(x,ch,3,3,1,1,name='sn_sh_upconv')
return x+short_cut
def optimized_block(x, out_channels, name, act=tf.nn.relu):
"""Builds optimized residual blocks for downsampling.
Compared with block, optimized_block always downsamples the spatial resolution
by a factor of 2 on each side.
Args:
x: The 4D input vector.
out_channels: Number of features in the output layer.
name: The variable scope name for the block.
act: The activation function used in the block.
Returns:
A `Tensor` representing the output of the operation.
"""
with tf.compat.v1.variable_scope(name):
x_0 = x
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv1')
x = act(x)
x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv2')
x = dsample(x)
x_0 = dsample(x_0)
x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='sn_conv3')
return x + x_0
def conv_block(inputs, num_filters, kernel_size, strides, padding, kernel_init, sn=True, norm_layer=None, activation=None):
if padding.lower() == 'same':
pad_size = (kernel_size - 1) // 2
inputs = tf.pad(inputs, [[0, 0], [pad_size, pad_size], [pad_size, pad_size], [0, 0]], mode='REFLECT')
else:
pad_size = 0
if sn:
out = snconv2d(inputs, num_filters, kernel_size, stride=strides, kernel_initializer=kernel_init)
else:
out = tf.layers.conv2d(inputs, num_filters, kernel_size, strides, padding, kernel_initializer=kernel_init)
if norm_layer is not None:
out = norm_layer(out)
if activation is not None:
out = activation(out)
return out
def generator(zs,
target_class,
gf_dim,
num_classes,
training=True,
mode='both'):
"""Builds the generator segment of the graph, going from z -> G(z).
Args:
zs: Tensor representing the latent variables.
target_class: The class from which we seek to sample.
gf_dim: The gf dimension.
num_classes: Number of classes in the labels.
training: Whether in train mode or not. This affects things like batch
normalization and spectral normalization.
Returns:
- The output layer of the generator.
- A list containing all trainable varaibles defined by the model.
"""
with tf.compat.v1.variable_scope(
'generator', reuse=tf.compat.v1.AUTO_REUSE) as gen_scope:
act0 = ops.snlinear(zs,
gf_dim * 16 * 4 * 4,
training=training,
name='g_snh0')
act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])
act1 = block(act0, target_class, gf_dim * 16, num_classes, 'g_block1',
training) # 8
act2 = block(act1, target_class, gf_dim * 8, num_classes, 'g_block2',
training) # 16
act3 = block(act2, target_class, gf_dim * 4, num_classes, 'g_block3',
training) # 32
act3, attn_map = ops.sn_attention_block_sim(act3,
training,
name='g_ops') # 32
act4 = block(act3, target_class, gf_dim * 2, num_classes, 'g_block4',
training) # 64
act5 = block(act4, target_class, gf_dim, num_classes, 'g_block5',
training) # 128
act5 = tf.nn.relu(
tfgan.tpu.batch_norm(act5,
training,
conditional_class_labels=None,
name='g_bn'))
act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, training, 'g_snconv_last')
out = tf.nn.tanh(act6)
return out, attn_map
def generator_test_64(zs,
target_class,
gf_dim,
num_classes,
CGN=False,
CGN_groups=4,
is_training=True):
"""Builds the generator graph propagating from z to x.
Args:
zs: The list of noise tensors.
target_class: The conditional labels in the generation.
gf_dim: The gf dimension.
num_classes: Number of classes in the labels.
scope: Optional scope for `variable_op_scope`.
Returns:
outputs: The output layer of the generator.
"""
with tf.variable_scope('model', reuse=tf.AUTO_REUSE):
# project `z` and reshape
act0 = ops.snlinear(zs, gf_dim * 16 * 4 * 4, name='g_snh0')
act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])
act1 = block(act0, target_class, gf_dim * 16, num_classes, is_training,
CGN, CGN_groups, 'g_block1') # 8 * 8
act2 = block(act1, target_class, gf_dim * 8, num_classes, is_training,
CGN, CGN_groups, 'g_block2') # 16 * 16
act3 = block(act2, target_class, gf_dim * 4, num_classes, is_training,
CGN, CGN_groups, 'g_block3') # 32 * 32
act4 = block(act3, target_class, gf_dim * 2, num_classes, is_training,
CGN, CGN_groups, 'g_block4') # 64 * 64
act4 = non_local.sn_non_local_block_sim(act4, None, name='g_non_local')
act5 = block(act4, target_class, gf_dim, num_classes, is_training, CGN,
CGN_groups, 'g_block5') # 128 * 128
bn = ops.batch_norm(name='g_bn')
act5 = tf.nn.relu(bn(act5, is_training))
act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, name='g_snconv_last')
out = tf.nn.tanh(act6)
print('GAN test with moving average')
return out
def __init__(self, z_size, channel, output_size=48):
super(Generator, self).__init__()
self.output_size = output_size
s = 4
if self.output_size == 48:
s = 6
self.s = s
self.z_size = z_size
self.fully_connect = snlinear(z_size, s * s * 512)
self.relu = nn.ReLU()
self.tanh = nn.Tanh()
self.conv_res4 = snconv2d(64,
channel,
padding=1,
kernel_size=3,
stride=1)
self.self_attn = Self_Attn(in_channels=256)
self.re1 = Residual_G(512, 256, up_sampling=True)
self.re2 = Residual_G(256, 128, up_sampling=True)
self.re3 = Residual_G(128, 64, up_sampling=True)
self.bn = nn.BatchNorm2d(64)
self.apply(init_weights)
