def build_C_global_layers(x, name, n_latents, start_idx, scope_idx, dlatents_withl_in, n_content, act, fused_modconv, fmaps=128, **kwargs): ''' Build continuous latent layers, e.g. C_global layers. ''' with tf.variable_scope(name + '-' + str(scope_idx)): if n_content > 0: with tf.variable_scope('Condition0'): cond = apply_bias_act(dense_layer( dlatents_withl_in[:, :n_content], fmaps=128), act=act) with tf.variable_scope('Condition1'): cond = apply_bias_act(dense_layer(cond, fmaps=n_latents), act='sigmoid') else: cond = 1. C_global_latents = dlatents_withl_in[:, start_idx:start_idx + n_latents] * cond x = apply_bias_act(modulated_conv2d_layer(x, C_global_latents, fmaps=fmaps, kernel=3, up=False, fused_modconv=fused_modconv), act=act) return x
def build_C_global_nocond_layers(x, name, n_latents, start_idx, scope_idx, dlatents_withl_in, act, fused_modconv, fmaps=128, **kwargs): ''' Build continuous latent layers, e.g. C_global layers. ''' with tf.variable_scope(name + '-' + str(scope_idx)): with tf.variable_scope('Conv0'): C_global_latents = apply_bias_act(dense_layer( dlatents_withl_in[:, start_idx:start_idx + n_latents], fmaps=128), act=act) # C_global_latents = dlatents_withl_in[:, start_idx:start_idx + # n_latents] with tf.variable_scope('Modulate'): x = apply_bias_act(modulated_conv2d_layer(x, C_global_latents, fmaps=fmaps, kernel=3, up=False, fused_modconv=fused_modconv), act=act) return x
def torgb(res, x): # res = 2..resolution_log2 with tf.variable_scope('ToRGB_lod%d' % (resolution_log2 - res)): return networks_stylegan2.apply_bias_act( networks_stylegan2.modulated_conv2d_layer( x, dlatents_in[:, res * 2 - 3], fmaps=num_channels, kernel=1, demodulate=False, fused_modconv=fused_modconv))
def torgb(x, y, res): # res = 2..resolution_log2 with tf.variable_scope('ToRGB'): t = apply_bias_act( modulated_conv2d_layer(x, latents_ready_ls[res * 2 - 3], fmaps=num_channels, kernel=1, demodulate=False, fused_modconv=fused_modconv)) return t if y is None else y + t
def build_D_layers(x, name, n_latents, start_idx, scope_idx, single_const, dlatents_withl_in, act, fused_modconv, fmaps=128, **kwargs): ''' Build discrete latent layers including label and D_global layers. ''' with tf.variable_scope(name + '-' + str(scope_idx)): if single_const: x = apply_bias_act(modulated_conv2d_layer( x, dlatents_withl_in[:, start_idx:start_idx + n_latents], fmaps=fmaps, kernel=3, up=False, fused_modconv=fused_modconv), act=act) else: # soft version # x_softmax = tf.nn.softmax( # dlatents_withl_in[:, start_idx:start_idx + n_latents], axis=1) # x_softmax = tf.reshape( # x_softmax, [tf.shape(x_softmax)[0], n_latents, 1, 1, 1]) # x = tf.reduce_sum(x * x_softmax, axis=1) # print('x.shape:', x.shape.as_list()) # hard version x_indices = tf.argmax(dlatents_withl_in[:, start_idx:start_idx + n_latents], axis=1) print('x_indices.shape:', x_indices.shape.as_list()) print('before gather_nd x.shape:', x.shape.as_list()) x = tf.gather(x, x_indices, axis=0) print('after gather_nd x.shape:', x.shape.as_list()) return x
def layer(x, layer_idx, fmaps, kernel, up=False): x = networks_stylegan2.modulated_conv2d_layer( x, dlatents_in[:, layer_idx], fmaps=fmaps, kernel=kernel, up=up, resample_kernel=resample_kernel, fused_modconv=fused_modconv) if layer_idx > num_layers * 0.75: if randomize_noise: noise = tf.random_normal( [tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype) else: noise = tf.cast(noise_inputs[layer_idx], x.dtype) noise_strength = tf.get_variable( 'noise_strength', shape=[], initializer=tf.initializers.zeros()) x += noise * tf.cast(noise_strength, x.dtype) return networks_stylegan2.apply_bias_act(x, act=act)
def layer(x, layer_idx, fmaps, kernel, up=False): # start_idx_layer = sum(latent_split_ls_for_std_gen[:layer_idx]) # for i in range(start_idx_layer, start_idx_layer + latent_split_ls_for_std_gen[layer_idx]): # x = modulated_conv2d_layer(x, latents_ready_spl_ls[i], fmaps=fmaps, kernel=kernel, up=up, # resample_kernel=resample_kernel, fused_modconv=fused_modconv) x = modulated_conv2d_layer(x, latents_ready_ls[layer_idx], fmaps=fmaps, kernel=kernel, up=up, resample_kernel=resample_kernel, fused_modconv=fused_modconv) if randomize_noise: noise = tf.random_normal( [tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype) else: noise = tf.cast(noise_inputs[layer_idx], x.dtype) noise_strength = tf.get_variable('noise_strength', shape=[], initializer=tf.initializers.zeros()) x += noise * tf.cast(noise_strength, x.dtype) return apply_bias_act(x, act=act)
def build_D_layers(x, name, n_latents, start_idx, scope_idx, dlatents_in, act, fused_modconv, fmaps=128, **kwargs): ''' Build discrete latent layers including label and D_global layers. ''' with tf.variable_scope(name + '-' + str(scope_idx)): x = apply_bias_act(modulated_conv2d_layer( x, dlatents_in[:, start_idx:start_idx + n_latents], fmaps=fmaps, kernel=3, up=False, fused_modconv=fused_modconv), act=act) return x
def G_synthesis_sb_general_dsp( dlatents_withl_in, # Input: Disentangled latents (W) [minibatch, label_size+dlatent_size]. dlatent_size=7, # Disentangled latent (W) dimensionality. Including discrete info, rotation, scaling, xy shearing, and xy translation. label_size=0, # Label dimensionality, 0 if no labels. D_global_size=3, # Global D_latents. C_global_size=0, # Global C_latents. sb_C_global_size=4, # Global spatial-biased C_latents. C_local_hfeat_size=0, # Local heatmap*features learned C_latents. C_local_heat_size=0, # Local heatmap learned C_latents. num_channels=1, # Number of output color channels. resolution=64, # Output resolution. nonlinearity='lrelu', # 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? use_noise=False, randomize_noise=True, # True = randomize noise inputs every time (non-deterministic), False = read noise inputs from variables. **_kwargs): # Ignore unrecognized keyword args. ''' dlatents_withl_in: dims contain: [label, D_global, C_global, sb_C_global, C_local_hfeat, C_local_feat] ''' resolution_log2 = int(np.log2(resolution)) # == 6 for resolution 64 assert resolution == 2**resolution_log2 and resolution >= 4 num_layers = resolution_log2 * 2 - 2 # == 10 for resolution 64 act = nonlinearity images_out = None # Primary inputs. assert dlatent_size == D_global_size + C_global_size + sb_C_global_size + \ C_local_hfeat_size + C_local_heat_size n_cat = label_size + D_global_size dlatents_withl_in.set_shape([None, label_size + dlatent_size]) dlatents_withl_in = tf.cast(dlatents_withl_in, dtype) n_content = label_size + D_global_size + C_global_size # Noise inputs. noise_inputs = [] for layer_idx in range(num_layers - 3): res = (layer_idx + 7) // 2 # [3, 4, 4, 5, 5, 6, 6] shape = [1, 1, 2**res, 2**res] noise_inputs.append( tf.get_variable('noise%d' % layer_idx, shape=shape, initializer=tf.initializers.random_normal(), trainable=False)) # Single convolution layer with all the bells and whistles. def noised_conv_layer(x, layer_idx, fmaps, kernel, up=False): x = conv2d_layer(x, fmaps=fmaps, up=up, kernel=kernel, resample_kernel=resample_kernel) if use_noise: if randomize_noise: noise = tf.random_normal( [tf.shape(x)[0], 1, x.shape[2], x.shape[3]], dtype=x.dtype) else: noise = tf.cast(noise_inputs[layer_idx], x.dtype) noise_strength = tf.get_variable( 'noise_strength', shape=[], initializer=tf.initializers.zeros()) x += noise * tf.cast(noise_strength, x.dtype) return apply_bias_act(x, act=act) # Early layers consists of 4x4 constant layer, # label+global discrete latents, # and global continuous latents. y = None with tf.variable_scope('4x4'): with tf.variable_scope('Const'): x = tf.get_variable('const', shape=[1, 128, 4, 4], initializer=tf.initializers.random_normal()) x = tf.tile(tf.cast(x, dtype), [tf.shape(dlatents_withl_in)[0], 1, 1, 1]) with tf.variable_scope('Upconv'): x = apply_bias_act(conv2d_layer(x, fmaps=128, kernel=3, up=True, resample_kernel=resample_kernel), act=act) with tf.variable_scope('8x8'): with tf.variable_scope('Conv0'): x = apply_bias_act(conv2d_layer(x, fmaps=128, kernel=3), act=act) with tf.variable_scope('Label_Dglobal_control'): x = apply_bias_act(modulated_conv2d_layer( x, dlatents_withl_in[:, :n_cat], fmaps=128, kernel=3, up=False, resample_kernel=resample_kernel, fused_modconv=fused_modconv), act=act) with tf.variable_scope('After_DiscreteGlobal_noised'): x = noised_conv_layer(x, layer_idx=0, fmaps=128, kernel=3) with tf.variable_scope('Cglobal_control'): start_idx = n_cat x = apply_bias_act(modulated_conv2d_layer( x, dlatents_withl_in[:, start_idx:start_idx + C_global_size], fmaps=128, kernel=3, up=False, resample_kernel=resample_kernel, fused_modconv=fused_modconv), act=act) with tf.variable_scope('After_ContinuousGlobal_noised'): x = noised_conv_layer(x, layer_idx=1, up=True, fmaps=128, kernel=3) # Spatial biased layers. with tf.variable_scope('16x16'): if C_local_hfeat_size > 0: with tf.variable_scope('LocalHFeat_C_latents'): with tf.variable_scope('ConstFeats'): const_feats = tf.get_variable( 'constfeats', shape=[1, C_local_hfeat_size, 32, 1, 1], initializer=tf.initializers.random_normal()) const_feats = tf.tile( tf.cast(const_feats, dtype), [tf.shape(const_feats)[0], 1, 1, 1, 1]) with tf.variable_scope('ControlAttHeat'): hfeat_start_idx = label_size + D_global_size + C_global_size + \ sb_C_global_size att_heat = get_att_heat(x, nheat=C_local_hfeat_size, act=act) att_heat = tf.reshape( att_heat, [tf.shape(att_heat)[0], C_local_hfeat_size, 1] + att_heat.shape.as_list()[2:4]) # C_local_heat latent [-2, 2] --> [0, 1] hfeat_modifier = (2 + dlatents_withl_in[:, hfeat_start_idx:hfeat_start_idx + \ C_local_hfeat_size]) / 4. hfeat_modifier = get_conditional_modifier( hfeat_modifier, dlatents_withl_in[:, :n_content], act=act) hfeat_modifier = tf.reshape( hfeat_modifier, [tf.shape(x)[0], C_local_hfeat_size, 1, 1, 1]) att_heat = att_heat * hfeat_modifier added_feats = const_feats * att_heat added_feats = tf.reshape(added_feats, [ tf.shape(att_heat)[0], C_local_hfeat_size * att_heat.shape.as_list()[2] ] + att_heat.shape.as_list()[3:5]) x = tf.concat([x, added_feats], axis=1) with tf.variable_scope('SpatialBiased_C_global'): # Rotation layers. start_idx = start_idx + C_global_size with tf.variable_scope('Rotation'): r_matrix = get_r_matrix( dlatents_withl_in[:, start_idx:start_idx + 1], dlatents_withl_in[:, :n_content], act=act) x = apply_st(x, r_matrix, up=False, fmaps=128, act=act) with tf.variable_scope('After_Rotation_noised'): x = noised_conv_layer(x, layer_idx=2, fmaps=128, kernel=3) # Scaling layers. start_idx = start_idx + 1 with tf.variable_scope('Scaling'): s_matrix = get_s_matrix( dlatents_withl_in[:, start_idx:start_idx + 1], dlatents_withl_in[:, :n_content], act=act) x = apply_st(x, s_matrix, up=False, fmaps=128, act=act) with tf.variable_scope('After_Scaling_noised'): x = noised_conv_layer(x, layer_idx=3, up=True, fmaps=128, kernel=3) with tf.variable_scope('32x32'): with tf.variable_scope('SpatialBiased_C_global'): # Shearing layers. with tf.variable_scope('Shearing'): start_idx = start_idx + 1 sh_matrix = get_sh_matrix( dlatents_withl_in[:, start_idx:start_idx + 2], dlatents_withl_in[:, :n_content], act=act) x = apply_st(x, sh_matrix, up=False, fmaps=128, act=act) with tf.variable_scope('After_Shearing_noised'): x = noised_conv_layer(x, layer_idx=4, fmaps=128, kernel=3) # Translation layers. with tf.variable_scope('Translation'): start_idx = start_idx + 2 t_matrix = get_t_matrix( dlatents_withl_in[:, start_idx:start_idx + 2], dlatents_withl_in[:, :n_content], act=act) x = apply_st(x, t_matrix, up=False, fmaps=128, act=act) with tf.variable_scope('After_Translation_noised'): if resolution_log2 >= 6: x = noised_conv_layer(x, layer_idx=5, up=True, fmaps=128, kernel=3) else: x = noised_conv_layer(x, layer_idx=5, fmaps=128, kernel=3) with tf.variable_scope('64x64' if resolution_log2 >= 6 else '32x32'): with tf.variable_scope('LocalHeat_C_latents'): with tf.variable_scope('ControlAttHeat'): heat_start_idx = label_size + D_global_size + C_global_size + \ sb_C_global_size + C_local_hfeat_size att_heat = get_att_heat(x, nheat=C_local_heat_size, act=act) # C_local_heat latent [-2, 2] --> [0, 1] heat_modifier = (2 + dlatents_withl_in[:, heat_start_idx:heat_start_idx + \ C_local_heat_size]) / 4. heat_modifier = get_conditional_modifier( heat_modifier, dlatents_withl_in[:, :n_content], act=act) heat_modifier = tf.reshape( heat_modifier, [tf.shape(heat_modifier)[0], C_local_heat_size, 1, 1]) att_heat = att_heat * heat_modifier x = tf.concat([x, att_heat], axis=1) with tf.variable_scope('After_LocalHeat_noised'): x = noised_conv_layer(x, layer_idx=6, fmaps=128, kernel=3) y = torgb(x, y, num_channels=num_channels) # # Tail layers. # for res in range(6, resolution_log2 + 1): # with tf.variable_scope('%dx%d' % (res * 2, res * 2)): # x = apply_bias_act(conv2d_layer(x, # fmaps=128, # kernel=1, # up=True, # resample_kernel=resample_kernel), # act=act) # y = torgb(x, y, num_channels=num_channels) images_out = y assert images_out.dtype == tf.as_dtype(dtype) return tf.identity(images_out, name='images_out')
def G_synthesis_spatial_biased_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. resolution=64, # Output resolution. fmap_base=16 << 10, # Overall multiplier for the number of feature maps. fmap_decay=1.0, # log2 feature map reduction when doubling the resolution. fmap_min=1, # Minimum number of feature maps in any layer. fmap_max=512, # Maximum number of feature maps in any layer. architecture='skip', # Architecture: 'orig', 'skip', 'resnet'. nonlinearity='lrelu', # 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. resolution_log2 = int(np.log2(resolution)) assert resolution == 2**resolution_log2 and resolution >= 4 def nf(stage): return np.clip(int(fmap_base / (2.0**(stage * fmap_decay))), fmap_min, fmap_max) assert architecture in ['orig', 'skip', 'resnet'] 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 shear matrix def get_sh_matrix(sh_latents, cond_latent): # sh_latents[:, 0]: [-2., 2.] -> [-1., 1.] # sh_latents[:, 1]: [-2., 2.] -> [-1., 1.] with tf.variable_scope('Condition0x'): cond_x = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act) with tf.variable_scope('Condition1x'): cond_x = apply_bias_act(dense_layer(cond_x, fmaps=1), act='sigmoid') with tf.variable_scope('Condition0y'): cond_y = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act) with tf.variable_scope('Condition1y'): cond_y = apply_bias_act(dense_layer(cond_y, fmaps=1), act='sigmoid') cond = tf.concat([cond_x, cond_y], axis=1) xy_shear = sh_latents / 2. * cond tt_00 = tf.ones_like(xy_shear[:, 0:1]) tt_01 = xy_shear[:, 0:1] tt_02 = tf.zeros_like(xy_shear[:, 0:1]) tt_10 = xy_shear[:, 1:] tt_11 = tf.ones_like(xy_shear[:, 1:]) tt_12 = tf.zeros_like(xy_shear[:, 1:]) 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('Condition0x'): cond_x = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act) with tf.variable_scope('Condition1x'): cond_x = apply_bias_act(dense_layer(cond_x, fmaps=1), act='sigmoid') with tf.variable_scope('Condition0y'): cond_y = apply_bias_act(dense_layer(cond_latent, fmaps=128), act=act) with tf.variable_scope('Condition1y'): cond_y = apply_bias_act(dense_layer(cond_y, fmaps=1), act='sigmoid') cond = tf.concat([cond_x, cond_y], axis=1) 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, idx, up=True): # idx: 2, 3, 4 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 with tf.variable_scope('Upconv'): x = apply_bias_act(conv2d_layer(x, fmaps=nf(idx), kernel=3, up=up, resample_kernel=resample_kernel), act=act) with tf.variable_scope('Conv'): x = apply_bias_act(conv2d_layer(x, fmaps=nf(idx), kernel=3), act=act) return x def upsample(y): with tf.variable_scope('Upsample'): return upsample_2d(y, k=resample_kernel) 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, nf(1), 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('Upconv8x8'): x = apply_bias_act(conv2d_layer(x, fmaps=nf(1), kernel=3, up=True, resample_kernel=resample_kernel), act=act) with tf.variable_scope('Conv0'): x = apply_bias_act(conv2d_layer(x, fmaps=nf(1), kernel=3), act=act) with tf.variable_scope('ModulatedConv'): x = apply_bias_act(modulated_conv2d_layer( x, dlatents_in[:, :n_cat], fmaps=nf(2), kernel=3, up=False, resample_kernel=resample_kernel, fused_modconv=fused_modconv), act=act) with tf.variable_scope('Conv1'): x = apply_bias_act(conv2d_layer(x, fmaps=nf(2), kernel=3), act=act) # Rotation layers. with tf.variable_scope('16x16'): r_matrix = get_r_matrix(dlatents_in[:, n_cat:n_cat + 1], dlatents_in[:, :n_cat]) x = apply_st(x, r_matrix, 2) # Scaling layers. with tf.variable_scope('32x32'): s_matrix = get_s_matrix(dlatents_in[:, n_cat + 1:n_cat + 2], dlatents_in[:, :n_cat]) x = apply_st(x, s_matrix, 3) # Shearing layers. with tf.variable_scope('32x32_Shear'): sh_matrix = get_sh_matrix(dlatents_in[:, n_cat + 2:n_cat + 4], dlatents_in[:, :n_cat]) x = apply_st(x, sh_matrix, 3, up=False) # Translation layers. with tf.variable_scope('64x64'): t_matrix = get_t_matrix(dlatents_in[:, n_cat + 4:], dlatents_in[:, :n_cat]) x = apply_st(x, t_matrix, 4) y = torgb(x, y) # # Tail layers. # for res in range(6, resolution_log2 + 1): # with tf.variable_scope('%dx%d' % (res * 2, res * 2)): # x = apply_bias_act(conv2d_layer(x, # fmaps=nf(res), # kernel=1, # up=True, # resample_kernel=resample_kernel), # act=act) # if architecture == 'skip': # y = upsample(y) # if architecture == 'skip' or res == resolution_log2: # y = torgb(x, y) images_out = y assert images_out.dtype == tf.as_dtype(dtype) return tf.identity(images_out, name='images_out')
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')