def forward(self, h, is_training): print(" [Build] Spatial Predictor ; is_training: {}".format(is_training)) update_collection = self._get_update_collection(is_training) with tf.variable_scope("Q_content_prediction_head", reuse=tf.AUTO_REUSE): h = snlinear(h, self.aux_dim, 'fc1', update_collection=update_collection) h = batch_norm(name='bn1')(h, is_training=is_training) h = lrelu(h) h = snlinear(h, self.z_dim, 'fc2', update_collection=update_collection) return tf.nn.tanh(h)
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 discriminator(image, labels, df_dim, number_classes, act=tf.nn.relu): """Builds the discriminator graph. Args: image: The current batch of images to classify as fake or real. labels: The corresponding labels for the images. df_dim: The df dimension. number_classes: The number of classes in the labels. act: The activation function used in the discriminator. Returns: - A `Tensor` representing the logits of the discriminator. - A list containing all trainable varaibles defined by the model. """ with tf.compat.v1.variable_scope( 'discriminator', reuse=tf.compat.v1.AUTO_REUSE) as dis_scope: h0 = optimized_block( image, df_dim, 'd_optimized_block1', act=act) # 64 * 64 h1_ = block(h0, df_dim * 2, 'd_block2', act=act) # 32 * 32 h1, attn_map = ops.sn_attention_block_sim(h1_, name='d_ops') # 32 * 32 h2 = block(h1, df_dim * 4, 'd_block3', act=act) # 16 * 16 h3 = block(h2, df_dim * 8, 'd_block4', act=act) # 8 * 8 h4 = block(h3, df_dim * 16, 'd_block5', act=act) # 4 * 4 h5 = block(h4, df_dim * 16, 'd_block6', downsample=False, act=act) h5_act = act(h5) h6 = tf.reduce_sum(input_tensor=h5_act, axis=[1, 2]) output = ops.snlinear(h6, 1, name='d_sn_linear') h_labels = ops.sn_embedding(labels, number_classes, df_dim * 16, name='d_embedding') output += tf.reduce_sum(input_tensor=h6 * h_labels, axis=1, keepdims=True) var_list = tf.compat.v1.get_collection( tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, dis_scope.name) return output, h1_, attn_map, var_list
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 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 __init__(self, ssup, channel, featOnly=False): super(Discriminator, self).__init__() self.ssup = ssup self.featOnly = featOnly self.lrelu = nn.LeakyReLU() self.relu = nn.ReLU() self.softmax = nn.Softmax() self.self_attn = Self_Attn(in_channels=64) self.re1 = Residual_D(channel, 64, down_sampling=True, is_start=True) self.re2 = Residual_D(64, 128, down_sampling=True) self.re3 = Residual_D(128, 256, down_sampling=True) self.re4 = Residual_D(256, 512, down_sampling=True) self.re5 = Residual_D(512, 1024) self.fully_connect_gan2 = snlinear(1024, 1) self.fully_connect_rot2 = snlinear(1024, 4) self.sigmoid = nn.Sigmoid() self.apply(init_weights)
def discriminator_test(image, labels, df_dim, number_classes, update_collection=None, act=tf.nn.relu): """Builds the discriminator graph. Args: image: The current batch of images to classify as fake or real. labels: The corresponding labels for the images. df_dim: The df dimension. number_classes: The number of classes in the labels. update_collection: The update collections used in the spectral_normed_weight. act: The activation function used in the discriminator. scope: Optional scope for `variable_op_scope`. Returns: A `Tensor` representing the logits of the discriminator. """ with tf.variable_scope('model', reuse=tf.AUTO_REUSE): h0 = optimized_block(image, df_dim, 'd_optimized_block1', update_collection, act=act) # 64 * 64 h1 = block(h0, df_dim * 2, 'd_block2', update_collection, act=act) # 32 * 32 h1 = non_local.sn_non_local_block_sim(h1, update_collection, name='d_non_local') # 32 * 32 h2 = block(h1, df_dim * 4, 'd_block3', update_collection, act=act) # 16 * 16 h3 = block(h2, df_dim * 8, 'd_block4', update_collection, act=act) # 8 * 8 h4 = block(h3, df_dim * 16, 'd_block5', update_collection, act=act) # 4 * 4 h5 = block(h4, df_dim * 16, 'd_block6', update_collection, False, act=act) h5_act = act(h5) h6 = tf.reduce_sum(h5_act, [1, 2]) output = ops.snlinear(h6, 1, update_collection=update_collection, name='d_sn_linear') h_labels = ops.sn_embedding(labels, number_classes, df_dim * 16, update_collection=update_collection, name='d_embedding') output += tf.reduce_sum(h6 * h_labels, axis=1, keepdims=True) print('Discriminator Test Structure') return output
def discriminator(image, labels, df_dim, number_classes, act=tf.nn.relu): """Builds the discriminator graph. Args: image: The current batch of images to classify as fake or real. labels: The corresponding labels for the images. df_dim: The df dimension. number_classes: The number of classes in the labels. act: The activation function used in the discriminator. Returns: - A `Tensor` representing the logits of the discriminator. - A list containing all trainable varaibles defined by the model. """ with tf.compat.v1.variable_scope( 'discriminator', reuse=tf.compat.v1.AUTO_REUSE) as dis_scope: h0 = optimized_block( image, df_dim, 'd_optimized_block1', act=act) # 64 * 64 h1 = block(h0, df_dim * 2, 'd_block2', act=act) # 32 * 32 if flags.FLAGS.D_module == 'hamburger': if flags.FLAGS.D_version == 'v1': print('Add V1 to D!') hamburger = ops.sn_hamburger_v1 else: print('Add V2 to D!') hamburger = ops.sn_hamburger_v2 h1 = hamburger(h1, name='d_ops', use_bn=False, ham_type=flags.FLAGS.D_ham_type, S=flags.FLAGS.D_s, D=flags.FLAGS.D_d, R=flags.FLAGS.D_r, steps=flags.FLAGS.D_K) else: print('No context module in D!') h2 = block(h1, df_dim * 4, 'd_block3', act=act) # 16 * 16 h3 = block(h2, df_dim * 8, 'd_block4', act=act) # 8 * 8 h4 = block(h3, df_dim * 16, 'd_block5', act=act) # 4 * 4 h5 = block(h4, df_dim * 16, 'd_block6', downsample=False, act=act) h5_act = act(h5) h6 = tf.reduce_sum(input_tensor=h5_act, axis=[1, 2]) output = ops.snlinear(h6, 1, name='d_sn_linear') h_labels = ops.sn_embedding(labels, number_classes, df_dim * 16, name='d_embedding') output += tf.reduce_sum(input_tensor=h6 * h_labels, axis=1, keepdims=True) var_list = tf.compat.v1.get_collection( tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, dis_scope.name) return output, var_list
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 discriminator_without_condition(image, df_dim, update_collection=None, act=tf.nn.relu): """Builds the discriminator graph. Args: image: The current batch of images to classify as fake or real. df_dim: The df dimension. update_collection: The update collections used in the spectral_normed_weight. act: The activation function used in the discriminator. scope: Optional scope for `variable_op_scope`. Returns: A `Tensor` representing the logits of the discriminator. """ with tf.variable_scope('model', reuse=tf.AUTO_REUSE): h0 = optimized_block(image, df_dim, 'd_optimized_block1', update_collection, act=act) # 64 * 64 h1 = block(h0, df_dim * 2, 'd_block2', update_collection, act=act) # 32 * 32 h1, attn = non_local.sn_non_local_block_sim( h1, update_collection, name='d_non_local') # 32 * 32 h2 = block(h1, df_dim * 4, 'd_block3', update_collection, act=act) # 16 * 16 h3 = block(h2, df_dim * 8, 'd_block4', update_collection, act=act) # 8 * 8 h4 = block(h3, df_dim * 16, 'd_block5', update_collection, act=act) # 4 * 4 h5 = block(h4, df_dim * 16, 'd_block6', update_collection, False, act=act) h5_act = act(h5) h6 = tf.reduce_sum(h5_act, [1, 2]) output = ops.snlinear(h6, 1, update_collection=update_collection, name='d_sn_linear') print('Discriminator without Condition with SA') return output, attn
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)
def forward(self, x, y=None, is_training=True): valid_sizes = {8, 16, 32, 64, 128, 256, 512} assert (self.macro_patch_size[0] in valid_sizes and self.macro_patch_size[1] in valid_sizes), \ "I haven't test your macro patch size: {}".format(self.macro_patch_size) update_collection = self._get_update_collection(is_training) print(" [Build] Discriminator ; is_training: {}".format(is_training)) with tf.variable_scope("D_discriminator", reuse=tf.AUTO_REUSE): num_resize_layers = int( math.log(min(self.macro_patch_size), 2) - 1) num_total_layers = num_resize_layers + self.num_extra_layers basic_layers = [2, 4, 8, 8] if num_total_layers > len(basic_layers): num_replicate_layers = num_total_layers - len(basic_layers) ndf_mult_list = [ 1, ] * num_replicate_layers + basic_layers else: ndf_mult_list = basic_layers[-num_total_layers:] ndf_mult_list[0] = 1 print("\t ndf_mult_list = {}".format(ndf_mult_list)) # Stack extra layers without resize first h = x for idx, ndf_mult in enumerate(ndf_mult_list): n_ch = self.ndf_base * ndf_mult # Head is fixed and goes first if idx == 0: is_head, resize, is_extra = True, True, False # Extra layers before standard layers elif idx <= self.num_extra_layers: is_head, resize, is_extra = False, False, True # Last standard layer has no resize elif idx == len(ndf_mult_list) - 1: is_head, resize, is_extra = False, False, False # Standard layers else: is_head, resize, is_extra = False, True, False h = self._d_residual_block(h, n_ch, idx=idx, is_training=is_training, resize=resize, is_head=is_head) print( "\t DResBlock: id={}, out_shape={}, resize={}, is_extra={}" .format(idx, h.shape.as_list(), resize, is_extra)) h = tf.nn.relu(h) h = tf.reduce_sum(h, axis=[1, 2]) # Global pooling last_feature_map = h adv_out = snlinear(h, 1, 'main_steam_out', update_collection=update_collection) # Projection Discriminator if y is not None: h_num_ch = self.ndf_base * ndf_mult_list[-1] y_emb = snlinear(y, h_num_ch, 'y_emb', update_collection=update_collection) proj_out = tf.reduce_sum(y_emb * h, axis=1, keepdims=True) else: proj_out = 0 out = adv_out + proj_out return out, last_feature_map