def dense_regW(x, fmaps, **kwargs):
if len(x.shape) > 2:
x = tf.reshape(x, [-1, np.prod([d.value for d in x.shape[1:]])])
# print('x.shape:', x.shape)
w = get_weight([x.shape[1].value, fmaps], **kwargs)
w = tf.cast(w, x.dtype)
return tf.matmul(x, w), w
def conv3d_layer(x,
fmaps,
kernel,
down=False,
resample_kernel=None,
gain=1,
use_wscale=True,
lrmul=1,
weight_var='weight'):
assert kernel >= 1 and kernel % 2 == 1
w = get_weight([3, kernel, kernel, x.shape[1].value, fmaps],
gain=gain,
use_wscale=use_wscale,
lrmul=lrmul,
weight_var=weight_var)
x = tf.nn.conv3d(x,
tf.cast(w, x.dtype),
data_format='NCDHW',
strides=[1, 1, 1, 1, 1],
padding='SAME')
if down:
x = tf.nn.avg_pool3d(x,
ksize=[1, 2, 2],
strides=[1, 2, 2],
padding='VALID',
data_format='NCDHW')
return x
def dense_layer_last_dim(x,
fmaps,
gain=1,
use_wscale=True,
lrmul=1,
weight_var='weight'):
w = get_weight([x.shape[-1].value, fmaps],
gain=gain,
use_wscale=use_wscale,
lrmul=lrmul,
weight_var=weight_var)
w = tf.cast(w, x.dtype)
return tf.matmul(x, w)
def net_M_hyperplane(
latents_in,
C_global_size=10,
D_global_size=0,
latent_size=512, # Latent vector (Z) dimensionality.
mapping_lrmul=0.1, # Learning rate multiplier for the mapping layers.
use_std_in_m=False, # If output prior std.
dtype='float32', # Data type to use for activations and outputs.
**_kwargs): # Ignore unrecognized keyword args.
latents_in.set_shape([None, C_global_size])
x = latents_in
# x = tf.concat([x, tf.zeros([tf.shape(x)[0], latent_size - C_global_size], dtype=x.dtype)], axis=1)
# W = tf.Variable(tf.random_normal((latent_size, latent_size)))
# W = get_weight([latent_size, latent_size], lrmul=mapping_lrmul, weight_var='orth_weight')
W = get_weight([C_global_size, latent_size],
lrmul=mapping_lrmul,
weight_var='weight')
W = tf.cast(W, x.dtype)
x = tf.matmul(x, W)
# # Orthogonal constraint on W
# constraint = tf.matmul(W, W, transpose_b=True) - tf.eye(latent_size, dtype=x.dtype)
# print('orth_constraint_1.shape:', constraint.get_shape().as_list())
# constraint = tf.reduce_sum(constraint * constraint)
# print('constraint.shape:', constraint.get_shape().as_list())
# Magnitude constraint on W
# constraint = tf.reduce_sum(W * W)
# W_norms = tf.math.sqrt(tf.reduce_sum(W * W, axis=1))
W_norms = tf.reduce_sum(W * W, axis=1)
constraint = W_norms - latent_size * tf.ones([C_global_size],
dtype=W_norms.dtype)
constraint = tf.reduce_mean(constraint * constraint)
# Output.
assert x.dtype == tf.as_dtype(dtype)
return tf.identity(x, name='to_latent_out'), constraint
def G_synthesis_simple_dsp(
dlatents_in, # Input: Disentangled latents (W) [minibatch, dlatent_size].
dlatent_size=7, # Disentangled latent (W) dimensionality. Including discrete info, rotation, scaling, and xy translation.
D_global_size=3, # Discrete latents.
sb_C_global_size=4, # Continuous latents.
label_size=0, # Label dimensionality, 0 if no labels.
num_channels=1, # Number of output color channels.
nonlinearity='relu', # Activation function: 'relu', 'lrelu', etc.
dtype='float32', # Data type to use for activations and outputs.
resample_kernel=[
1, 3, 3, 1
], # Low-pass filter to apply when resampling activations. None = no filtering.
fused_modconv=True, # Implement modulated_conv2d_layer() as a single fused op?
**_kwargs): # Ignore unrecognized keyword args.
act = nonlinearity
images_out = None
# Primary inputs.
assert dlatent_size == D_global_size + sb_C_global_size
n_cat = label_size + D_global_size
dlatents_in.set_shape([None, label_size + dlatent_size])
dlatents_in = tf.cast(dlatents_in, dtype)
# Return rotation matrix
def get_r_matrix(r_latents, cond_latent):
# r_latents: [-2., 2.] -> [0, 2*pi]
with tf.variable_scope('Condition0'):
cond = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act)
with tf.variable_scope('Condition1'):
cond = apply_bias_act(dense_layer(cond, fmaps=1), act='sigmoid')
rad = (r_latents + 2) / 4. * 2. * np.pi
rad = rad * cond
tt_00 = tf.math.cos(rad)
tt_01 = -tf.math.sin(rad)
tt_02 = tf.zeros_like(rad)
tt_10 = tf.math.sin(rad)
tt_11 = tf.math.cos(rad)
tt_12 = tf.zeros_like(rad)
theta = tf.concat([tt_00, tt_01, tt_02, tt_10, tt_11, tt_12], axis=1)
return theta
# Return scaling matrix
def get_s_matrix(s_latents, cond_latent):
# s_latents: [-2., 2.] -> [1, 3]
with tf.variable_scope('Condition0'):
cond = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act)
with tf.variable_scope('Condition1'):
cond = apply_bias_act(dense_layer(cond, fmaps=1), act='sigmoid')
scale = (s_latents / 2. + 2.) * cond
tt_00 = scale
tt_01 = tf.zeros_like(scale)
tt_02 = tf.zeros_like(scale)
tt_10 = tf.zeros_like(scale)
tt_11 = scale
tt_12 = tf.zeros_like(scale)
theta = tf.concat([tt_00, tt_01, tt_02, tt_10, tt_11, tt_12], axis=1)
return theta
# Return translation matrix
def get_t_matrix(t_latents, cond_latent):
# t_latents[:, 0]: [-2., 2.] -> [-0.5, 0.5]
# t_latents[:, 1]: [-2., 2.] -> [-0.5, 0.5]
with tf.variable_scope('Condition0'):
cond = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act)
with tf.variable_scope('Condition1'):
cond = apply_bias_act(dense_layer(cond, fmaps=1), act='sigmoid')
xy_shift = t_latents / 4. * cond
tt_00 = tf.ones_like(xy_shift[:, 0:1])
tt_01 = tf.zeros_like(xy_shift[:, 0:1])
tt_02 = xy_shift[:, 0:1]
tt_10 = tf.zeros_like(xy_shift[:, 1:])
tt_11 = tf.ones_like(xy_shift[:, 1:])
tt_12 = xy_shift[:, 1:]
theta = tf.concat([tt_00, tt_01, tt_02, tt_10, tt_11, tt_12], axis=1)
return theta
# Apply spatial transform
def apply_st(x, st_matrix):
with tf.variable_scope('Transform'):
x = tf.transpose(x, [0, 2, 3, 1]) # NCHW -> NHWC
x = transformer(x, st_matrix, out_dims=x.shape.as_list()[1:3])
x = tf.transpose(x, [0, 3, 1, 2]) # NHWC -> NCHW
return x
def torgb(x, y):
with tf.variable_scope('ToRGB'):
t = apply_bias_act(conv2d_layer(x, fmaps=num_channels, kernel=1))
return t if y is None else y + t
# Early layers.
y = None
with tf.variable_scope('4x4'):
with tf.variable_scope('Const'):
x = tf.get_variable('const',
shape=[1, 64, 4, 4],
initializer=tf.initializers.random_normal())
x = tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_in)[0], 1, 1, 1])
with tf.variable_scope('4x4Conv'):
w = get_weight([3, 3, x.shape[1].value, 64])
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 1, 1],
padding='SAME')
x = apply_bias_act(x, act=act)
with tf.variable_scope('8x8ModulatedConv'):
x = apply_bias_act(modulated_conv2d_layer(
x,
dlatents_in[:, :n_cat],
fmaps=64,
kernel=3,
up=True,
resample_kernel=resample_kernel,
fused_modconv=fused_modconv),
act=act)
with tf.variable_scope('16x16'):
w = get_weight([4, 4, x.shape[1].value, 32])
# Transpose weights.
w = tf.transpose(w, [0, 1, 3, 2])
x = tf.nn.conv2d_transpose(
x,
w,
output_shape=[tf.shape(dlatents_in)[0], 32, 16, 16],
strides=[1, 1, 2, 2],
padding='SAME',
data_format='NCHW')
x = apply_bias_act(x, act=act)
with tf.variable_scope('rotation'):
r_matrix = get_r_matrix(dlatents_in[:, n_cat:n_cat + 1],
dlatents_in[:, :n_cat])
x = apply_st(x, r_matrix)
with tf.variable_scope('scale'):
s_matrix = get_s_matrix(dlatents_in[:, n_cat + 1:n_cat + 2],
dlatents_in[:, :n_cat])
x = apply_st(x, s_matrix)
with tf.variable_scope('translation'):
t_matrix = get_t_matrix(dlatents_in[:, n_cat + 2:],
dlatents_in[:, :n_cat])
x = apply_st(x, t_matrix)
with tf.variable_scope('32x32'):
w = get_weight([4, 4, x.shape[1].value, 32])
# Transpose weights.
w = tf.transpose(w, [0, 1, 3, 2])
x = tf.nn.conv2d_transpose(
x,
w,
output_shape=[tf.shape(dlatents_in)[0], 32, 32, 32],
strides=[1, 1, 2, 2],
padding='SAME',
data_format='NCHW')
x = apply_bias_act(x, act=act)
with tf.variable_scope('64x64'):
w = get_weight([4, 4, x.shape[1].value, 1])
# Transpose weights.
w = tf.transpose(w, [0, 1, 3, 2])
x = tf.nn.conv2d_transpose(
x,
w,
output_shape=[tf.shape(dlatents_in)[0], 1, 64, 64],
strides=[1, 1, 2, 2],
padding='SAME',
data_format='NCHW')
x = apply_bias_act(x)
images_out = x
assert images_out.dtype == tf.as_dtype(dtype)
return tf.identity(images_out, name='images_out')
def D_simple_dsp(
images_in, # First input: Images [minibatch, channel, height, width].
labels_in, # Second input: Labels [minibatch, label_size].
num_channels=1, # Number of input color channels. Overridden based on dataset.
label_size=0, # Dimensionality of the labels, 0 if no labels. Overridden based on dataset.
nonlinearity='relu', # Activation function: 'relu', 'lrelu', etc.
dtype='float32', # Data type to use for activations and outputs.
**_kwargs): # Ignore unrecognized keyword args.
act = nonlinearity
images_in.set_shape([None, num_channels, 64, 64])
labels_in.set_shape([None, label_size])
images_in = tf.cast(images_in, dtype)
labels_in = tf.cast(labels_in, dtype)
x = images_in
with tf.variable_scope('32x32'):
w = get_weight([4, 4, x.shape[1].value, 32])
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 2, 2],
padding='SAME')
x = apply_bias_act(x, act=act)
with tf.variable_scope('16x16'):
w = get_weight([4, 4, x.shape[1].value, 32])
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 2, 2],
padding='SAME')
x = apply_bias_act(x, act=act)
with tf.variable_scope('8x8'):
w = get_weight([4, 4, x.shape[1].value, 64])
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 2, 2],
padding='SAME')
x = apply_bias_act(x, act=act)
with tf.variable_scope('4x4'):
w = get_weight([4, 4, x.shape[1].value, 64])
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 2, 2],
padding='SAME')
x = apply_bias_act(x, act=act)
with tf.variable_scope('output'):
x = tf.reshape(x, [-1, np.prod([d.value for d in x.shape[1:]])])
w = get_weight([x.shape[1].value, 128])
w = tf.cast(w, x.dtype)
x = tf.matmul(x, w)
x = apply_bias_act(x, act=act)
# Output layer with label conditioning from "Which Training Methods for GANs do actually Converge?"
with tf.variable_scope('Score'):
w = get_weight([x.shape[1].value, 1])
w = tf.cast(w, x.dtype)
x = tf.matmul(x, w)
x = apply_bias_act(x)
scores_out = x
# Output.
assert scores_out.dtype == tf.as_dtype(dtype)
scores_out = tf.identity(scores_out, name='scores_out')
return scores_out