def block(x, res): # res = 2..resolution_log2 t = x with tf.variable_scope('Conv0'): x = apply_bias_act(modulated_conv2d_layer(x, dlabel, fmaps=nf(res - 1), kernel=3), act=act) with tf.variable_scope('Conv1_down'): x = apply_bias_act(modulated_conv2d_layer( x, dlabel, fmaps=nf(res - 2), kernel=3, down=True, resample_kernel=resample_kernel), act=act) if architecture == 'resnet': with tf.variable_scope('Skip'): t = conv2d_layer(t, fmaps=nf(res - 2), kernel=1, down=True, resample_kernel=resample_kernel) x = (x + t) * (1 / np.sqrt(2)) return x
def fromrgb(x, y, res): # res = 2..resolution_log2 with tf.variable_scope('FromRGB'): t = apply_bias_act(modulated_conv2d_layer(y, dlabel, fmaps=nf(res - 1), kernel=1), act=act) return t if x is None else x + t
def layer(x, layer_idx, fmaps, kernel, up=False): x = modulated_conv2d_layer(x, dlatents_in[:, 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 torgb(res, x): # res = 2..resolution_log2 with tf.variable_scope('ToRGB_lod%d' % (resolution_log2 - res)): return apply_bias_act( 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, dlatents_in[:, res * 2 - 3], fmaps=num_channels, kernel=1, demodulate=False, fused_modconv=fused_modconv)) return t if y is None else y + t
def D_stylegan2( images_in, # First input: Images [minibatch, channel, height, width]. labels_in, # Second input: Labels [minibatch, label_size]. num_channels = 3, # Number of input color channels. Overridden based on dataset. resolution = 512, # Input resolution. Overridden based on dataset. label_size = 127, # Dimensionality of the labels, 0 if no labels. Overridden based on dataset. 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 = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu', etc. mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, 0 = disable. mbstd_num_features = 1, # Number of features for the minibatch standard deviation layer. 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. **_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_in.set_shape([None, num_channels, resolution, resolution]) labels_in.set_shape([None, label_size]) images_in = tf.cast(images_in, dtype) labels_in = tf.cast(labels_in, dtype) dlabel = D_mapping_label(labels_in=labels_in, label_size=label_size) dlabel = tf.cast(dlabel, dtype) # Building blocks for main layers. def fromrgb(x, y, res): # res = 2..resolution_log2 with tf.variable_scope('FromRGB'): t = apply_bias_act(modulated_conv2d_layer(y, dlabel, fmaps=nf(res-1), kernel=1), act=act) return t if x is None else x + t def block(x, res): # res = 2..resolution_log2 t = x with tf.variable_scope('Conv0'): x = apply_bias_act(modulated_conv2d_layer(x, dlabel, fmaps=nf(res-1), kernel=3), act=act) with tf.variable_scope('Conv1_down'): x = apply_bias_act(modulated_conv2d_layer(x, dlabel, fmaps=nf(res-2), kernel=3, down=True, resample_kernel=resample_kernel), act=act) if architecture == 'resnet': with tf.variable_scope('Skip'): t = conv2d_layer(t, fmaps=nf(res-2), kernel=1, down=True, resample_kernel=resample_kernel) x = (x + t) * (1 / np.sqrt(2)) return x def downsample(y): with tf.variable_scope('Downsample'): return downsample_2d(y, k=resample_kernel) feature_maps = [] # Main layers. x = None y = images_in for res in range(resolution_log2, 2, -1): with tf.variable_scope('%dx%d' % (2**res, 2**res)): if architecture == 'skip' or res == resolution_log2: x = fromrgb(x, y, res) feature_maps.append(x) x = block(x, res) if architecture == 'skip': y = downsample(y) feature_maps.append(x) # Final layers. with tf.variable_scope('4x4'): if architecture == 'skip': x = fromrgb(x, y, 2) if mbstd_group_size > 1: with tf.variable_scope('MinibatchStddev'): x = minibatch_stddev_layer(x, mbstd_group_size, mbstd_num_features) with tf.variable_scope('Conv'): x = apply_bias_act(modulated_conv2d_layer(x, dlabel, fmaps=nf(1), kernel=3), act=act) with tf.variable_scope('Dense0'): x = apply_bias_act(dense_layer(x, fmaps=nf(0)), act=act) # Output layer with label conditioning from "Which Training Methods for GANs do actually Converge?" with tf.variable_scope('Output'): x = apply_bias_act(dense_layer(x, fmaps=max(labels_in.shape[1], 1))) if labels_in.shape[1] > 0: x = tf.reduce_sum(x * labels_in, axis=1, keepdims=True) scores_out = x # Output. assert scores_out.dtype == tf.as_dtype(dtype) scores_out = tf.identity(scores_out, name='scores_out') feature_maps_0 = tf.identity(feature_maps[0], name='feature_maps_0') feature_maps_1 = tf.identity(feature_maps[1], name='feature_maps_1') feature_maps_2 = tf.identity(feature_maps[2], name='feature_maps_2') feature_maps_3 = tf.identity(feature_maps[3], name='feature_maps_3') feature_maps_4 = tf.identity(feature_maps[4], name='feature_maps_4') feature_maps_5 = tf.identity(feature_maps[5], name='feature_maps_5') feature_maps_6 = tf.identity(feature_maps[6], name='feature_maps_6') feature_maps_7 = tf.identity(feature_maps[7], name='feature_maps_7') return scores_out, feature_maps_0, feature_maps_1, feature_maps_2, feature_maps_3, feature_maps_4, feature_maps_5, feature_maps_6, feature_maps_7