def discriminator(x, nf=64): ''' :param x: input to the network :param nf: number of output channels :return: ''' # [3,128, 128] h = F.leaky_relu(PF.convolution(x, nf, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv0_0', with_bias=True), alpha=0.2) # [64, 128, 128] h = conv_bn_4(h, nf, "conv0_1", False) h = conv_bn_3(h, 2 * nf, "conv1_0", False) h = conv_bn_4(h, 2 * nf, "conv1_1", False) h = conv_bn_3(h, 4 * nf, "conv2_0", False) h = conv_bn_4(h, 4 * nf, "conv2_1", False) h = conv_bn_3(h, 8 * nf, "conv3_0", False) h = conv_bn_4(h, 8 * nf, "conv3_1", False) h = conv_bn_3(h, 8 * nf, "conv4_0", False) h = conv_bn_4(h, 8 * nf, "conv4_1", False) # [512, 4, 4] B, C, H, W = h.shape[0], h.shape[1], h.shape[2], h.shape[3] h = F.leaky_relu((PF.affine(h, 100, name="affine1")), alpha=0.2) h = PF.affine(h, 1, name="affine2") return h
def call(self, x, p): r"""Returns discriminator period. Args: x (nn.Variable): Input variable of shape (B, 1, L). p (int): Period. Returns: List[nn.Variable]: List of feature maps. """ results = list() b, c, t = x.shape if t % p: x = F.pad(x, (0, 0, 0, p - (t % p)), 'reflect') t = x.shape[-1] x = F.reshape(x, (b, c, t // p, p)) for i, c in enumerate([32, 128, 512, 1024, 1024]): with nn.parameter_scope(f"conv_{i}"): x = wn_conv(x, c, (5, 1), stride=(3, 1) if i < 4 else (1, 1), pad=(2, 0)) x = F.leaky_relu(x, 0.1) results.append(x) with nn.parameter_scope("post_conv"): x = wn_conv(x, 1, (3, 1), pad=(1, 0)) x = F.leaky_relu(x, 0.1) results.append(x) return results
def dis_block(n, c, i, train=True): """ Discriminator conv_bn_relu block """ out = conv(n, channels=c, kernel=4, stride=2, pad=1, use_bias=False, scope='d_conv/' + str(2 * i + 2)) out_fm = F.leaky_relu(PF.batch_normalization(out, axes=[3], batch_stat=train, name='d_bn/' + str(2 * i + 1)), alpha=0.2) out = conv(out_fm, channels=c * 2, kernel=3, stride=1, pad=1, use_bias=False, scope='d_conv/' + str(2 * i + 3)) out = F.leaky_relu(PF.batch_normalization(out, axes=[3], batch_stat=train, name='d_bn/' + str(2 * i + 2)), alpha=0.2) return out, out_fm
def DownsampleComp(h, nmap_out, scope_name): with nn.parameter_scope(scope_name): def sn_w(w): return PF.spectral_norm(w, dim=0) # Main h0 = PF.convolution(h, nmap_out, (4, 4), stride=(2, 2), pad=(1, 1), apply_w=sn_w, with_bias=False, name="main_conv1") h0 = PF.batch_normalization(h0, name="bn_main1") h0 = F.leaky_relu(h0, 0.2) h0 = PF.convolution(h0, nmap_out, (3, 3), pad=(1, 1), apply_w=sn_w, with_bias=False, name="main_conv2") h0 = PF.batch_normalization(h0, name="bn_main2") h0 = F.leaky_relu(h0, 0.2) # Direct h1 = F.average_pooling(h, (2, 2), stride=(2, 2)) h1 = PF.convolution(h1, nmap_out, (1, 1), apply_w=sn_w, with_bias=False, name="direct_conv1") h1 = PF.batch_normalization(h1, name="direct_bn1") h1 = F.leaky_relu(h1, 0.2) return (h0 + h1) / 2.0
def res_block(x, dim_out, kernel, pad, scope, training): r"""Residual block. Args: x (nn.Variable): Input variable. dim_out (int): Number of output channels. kernel (tuple of int): Kernel size. pad (tuple of int): Padding. scope (str): The scope. Returns: nn.Variable: Output variable. """ dim_in = x.shape[1] with nn.parameter_scope(scope): with nn.parameter_scope('shortcut'): sc = x if dim_in != dim_out: sc = wn_conv(sc, dim_out, (1, ), with_bias=False, name='conv') sc = F.average_pooling(sc, kernel=(1, 2)) with nn.parameter_scope('residual'): x = F.leaky_relu(x, 0.2) x = wn_conv(x, dim_in, kernel, pad, name='conv1') x = F.average_pooling(x, kernel=(1, 2)) x = F.leaky_relu(x, 0.2, inplace=True) x = wn_conv(x, dim_out, (1, ), name='conv2') return sc + x
def network(x, y, test=False): # Input:x -> 3,64,64 # AveragePooling -> 3,12,21 h = F.average_pooling(x, (5, 3), (5, 3)) # LeakyReLU_2 h = F.leaky_relu(h, 0.1, True) # Convolution_2 -> 20,13,21 h = PF.convolution(h, 20, (2, 3), (1, 1), name='Convolution_2') # BatchNormalization h = PF.batch_normalization(h, (1, ), 0.9, 0.0001, not test, name='BatchNormalization') # ReLU h = F.relu(h, True) # DepthwiseConvolution h = PF.depthwise_convolution(h, (5, 5), (2, 2), name='DepthwiseConvolution') # MaxPooling_2 -> 20,6,7 h = F.max_pooling(h, (2, 3), (2, 3)) # LeakyReLU h = F.leaky_relu(h, 0.1, True) # Affine -> 2 h = PF.affine(h, (2, ), name='Affine') # Softmax h = F.softmax(h) return h
def G(z, is_train=True): z = F.reshape(z, [batch_size, z_dim, 1, 1]) with nn.parameter_scope('G'): with nn.parameter_scope('deconv1'): dc1 = PF.deconvolution(z, 256, (4, 4), with_bias=False) dc1 = PF.batch_normalization(dc1, batch_stat=is_train) dc1 = F.leaky_relu(dc1) with nn.parameter_scope('deconv2'): dc2 = PF.deconvolution(dc1, 128, (4, 4), pad=(1, 1), stride=(2, 2), with_bias=False) dc2 = PF.batch_normalization(dc2, batch_stat=is_train) dc2 = F.leaky_relu(dc2) with nn.parameter_scope('deconv3'): dc3 = PF.deconvolution(dc2, 64, (4, 4), pad=(1, 1), stride=(2, 2), with_bias=False) dc3 = PF.batch_normalization(dc3, batch_stat=is_train) dc3 = F.leaky_relu(dc3) with nn.parameter_scope('deconv4'): dc4 = PF.deconvolution(dc3, 32, (4, 4), pad=(3, 3), stride=(2, 2), with_bias=False) dc4 = PF.batch_normalization(dc4, batch_stat=is_train) dc4 = F.leaky_relu(dc4) with nn.parameter_scope('output'): output = PF.convolution(dc4, 1, (3, 3), pad=(1, 1)) output = F.sigmoid(output) return output
def rrdb_net(x, num_output_channel, num_rrdb_blocks, growth_channel=32): ''' :param x: input image :param num_output_channel: number of output channels :param num_rrdb_blocks: number of residual blocks :param growth_channel: growth channel (no. of intermediate channel) :return: ''' fea = PF.convolution(x, num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv_first') h = fea with nn.parameter_scope('RRDB_trunk'): for i in range(num_rrdb_blocks): with nn.parameter_scope('{}'.format(i)): h = rrdb(h, num_output_channel, growth_channel) trunk_conv = PF.convolution(h, num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='trunk_conv') fea = fea + trunk_conv up_conv1 = F.leaky_relu(PF.convolution(F.interpolate(fea, scale=(2, 2), mode='nearest'), num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='upconv1'), alpha=0.2) up_conv2 = F.leaky_relu(PF.convolution(F.interpolate(up_conv1, scale=(2, 2), mode='nearest'), num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='upconv2'), alpha=0.2) hr_conv = F.leaky_relu(PF.convolution(up_conv2, num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='HRconv'), alpha=0.2) conv_last = PF.convolution(hr_conv, 3, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv_last') return conv_last
def down_block(input, output_channels=64, stride=1, scope='down_block'): with nn.parameter_scope(scope): net = conv2d(input, output_channels, (3, 3), (stride, stride), name='conv_1') net = F.leaky_relu(net, 0.2) net = conv2d(net, output_channels, (3, 3), (stride, stride), name='conv_2') net = F.leaky_relu(net, 0.2) net = F.max_pooling(net, (2, 2), channel_last=True) return net
def up_block(input, output_channels=64, stride=1, scope='up_block'): with nn.parameter_scope(scope): net = conv2d(input, output_channels, (3, 3), (stride, stride), name='conv_1') net = F.leaky_relu(net, 0.2) net = conv2d(net, output_channels, (3, 3), (stride, stride), name='conv_2') net = F.leaky_relu(net, 0.2) net = F.interpolate(net, scale=(2, 2), channel_last=True) return net
def discriminator(x): with nn.parameter_scope("discriminator"): with nn.parameter_scope("block1"): x = pfunc.affine(x, 1024) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block2"): x = pfunc.affine(x, 512) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block3"): x = pfunc.affine(x, 256) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block4"): x = pfunc.affine(x, 1) x = func.sigmoid(x) return x
def generator(x): with nn.parameter_scope("generator"): with nn.parameter_scope("block1"): x = pfunc.affine(x, 256) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block2"): x = pfunc.affine(x, 512) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block3"): x = pfunc.affine(x, 1024) x = func.leaky_relu(x, 0.2) with nn.parameter_scope("block4"): x = pfunc.affine(x, 784) x = func.tanh(x) return x
def conv_bn_4(x, nf, name, bias): with nn.parameter_scope(name): h = PF.convolution(x, nf, kernel=(4, 4), stride=( 2, 2), pad=(1, 1), with_bias=bias) h = PF.batch_normalization(h) h = F.leaky_relu(h, alpha=0.2) return h
def downblock(x, out_features, norm=False, kernel_size=4, pool=False, sn=False, test=False): out = x if sn: def apply_w(w): return PF.spectral_norm(w, dim=0, test=test) else: apply_w = None inmaps, outmaps = out.shape[1], out_features k_w = I.calc_normal_std_he_forward( inmaps, outmaps, kernel=(kernel_size, kernel_size)) / np.sqrt(2.) k_b = I.calc_normal_std_he_forward(inmaps, outmaps) / np.sqrt(2.) w_init = I.UniformInitializer((-k_w, k_w)) b_init = I.UniformInitializer((-k_b, k_b)) out = PF.convolution(out, out_features, kernel=(kernel_size, kernel_size), pad=(0, 0), stride=(1, 1), w_init=w_init, b_init=b_init, apply_w=apply_w) if norm: out = PF.instance_normalization(out) out = F.leaky_relu(out, 0.2, inplace=True) if pool: out = F.average_pooling(out, kernel=(2, 2)) return out
def unpool_block(x, n=0, k=(4, 4), s=(2, 2), p=(1, 1), leaky=False, unpool=False, init_method=None): if not unpool: logger.info("Deconvolution was used.") x = deconvolution(x, n=n, kernel=k, stride=s, pad=p, init_method=init_method) else: logger.info("Unpooling was used.") x = F.unpooling(x, kernel=(2, 2)) x = convolution(x, n, kernel=(3, 3), stride=(1, 1), pad=(1, 1), init_method=init_method) x = instance_normalization(x, fix_parameters=True) x = F.leaky_relu(x, alpha=0.2) if leaky else F.relu(x) return x
def call(self, x): hp = self.hp kernel, pad = (3, ), (1, ) with nn.parameter_scope('melspectrogram'): out = log_mel_spectrogram(x, hp.sr, 1024) with nn.parameter_scope('init_conv'): out = wn_conv(out, 32, kernel=kernel, pad=pad) for i in range(5): dim_out = min(out.shape[1] * 2, 512) out = res_block(out, dim_out, kernel=kernel, pad=pad, scope=f'downsample_{i}', training=self.training) with nn.parameter_scope('last_layer'): out = F.mean(out, axis=(2, )) out = F.leaky_relu(out, 0.2) out = PF.affine(out, hp.n_speakers) return out
def call(self, x, i): hp = self.hp dilations = hp.resblock_dilation_sizes kernel_sizes = hp.resblock_kernel_sizes res_block = res_block_1 if hp.resblock == "1" else res_block_2 channel = hp.upsample_initial_channel // (2**(i + 1)) r = hp.upsample_rates[i] with nn.parameter_scope("deconv"): x = F.leaky_relu(x, 0.1) x = wn_deconv( x, channel, (r * 2, ), stride=(r, ), pad=(r // 2 + r % 2, ), ) out = list() for j, (k, d) in enumerate(zip(kernel_sizes, dilations)): with nn.parameter_scope(f"resblock_{j}"): out.append(res_block(x, channel, k, d)) x = sum(out) / len(kernel_sizes) return x
def discriminator(x, scopename="discriminator", maps=128): with nn.parameter_scope(scopename): h = PF.convolution(x, maps * 1, (3, 3), name="conv-1") h = F.average_pooling(h, (2, 2)) # 32x32 -> 16x16 h = F.leaky_relu(h, 0.01) h = PF.convolution(h, maps * 2, (3, 3), name="conv-2") h = F.average_pooling(h, (2, 2)) # 16x16 -> 8x8 h = F.leaky_relu(h, 0.01) h = PF.convolution(h, maps * 4, (3, 3), name="conv-3") h = F.average_pooling(h, (2, 2)) # 8x8 -> 4x4 h = F.leaky_relu(h, 0.01) h = PF.affine(h, 1) return h
def Downsample(h, nmap_out, scope_name): with nn.parameter_scope(scope_name): def sn_w(w): return PF.spectral_norm(w, dim=0) h = PF.convolution(h, nmap_out, (4, 4), stride=( 2, 2), pad=(1, 1), apply_w=sn_w, with_bias=False) h = PF.batch_normalization(h) h = F.leaky_relu(h, 0.2, inplace=True) return h
def conv_block(x, out_ch, kernel, stride, pad, test, scope): with nn.parameter_scope(scope): h = PF.convolution(x, out_ch, (kernel, kernel), stride=(stride, stride), pad=(pad, pad), w_init=NormalInitializer(0.02), name='conv') h = PF.batch_normalization(h, batch_stat=not test, name='bn') h = F.leaky_relu(h, alpha=0.2) return h
def __call__(self, z, m): # m has target image shape: (N, emb, H, W) # z: (N, z_dim) N = m.shape[0] H, W = self.image_shape sh = H // (2**self.num_upsample) sw = W // (2**self.num_upsample) with ps("spade_generator"): with ps("z_embedding"): x = PF.affine(z, 16 * self.nf * sh * sw, w_init=w_init(z, 16 * self.nf * sh * sw)) x = F.reshape(x, (N, 16 * self.nf, sh, sw)) with ps("head"): x = self.head_0(x, m) with ps("middle0"): x = self.up(x) x = self.G_middle_0(x, m) with ps("middel1"): if self.num_upsample > 5: x = self.up(x) x = self.G_middle_1(x, m) with ps("up0"): x = self.up(x) x = self.up_0(x, m) with ps("up1"): x = self.up(x) x = self.up_1(x, m) with ps("up2"): x = self.up(x) x = self.up_2(x, m) with ps("up3"): x = self.up(x) x = self.up_3(x, m) if self.num_upsample > 6: with ps("up4"): x = self.up(x) x = self.up_4(x, m) with ps("last_conv"): x = PF.convolution(F.leaky_relu(x, 2e-1), 3, kernel=(3, 3), pad=(1, 1), w_init=w_init(x, 3)) x = F.tanh(x) return x
def residual_block_5C(x, num_output_channel=64, growth_channel=32): conv1 = F.leaky_relu(PF.convolution(x, growth_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv1'), alpha=0.2, inplace=True) conv2 = F.leaky_relu(PF.convolution(F.concatenate(x, conv1, axis=1), growth_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv2'), alpha=0.2, inplace=True) conv3 = F.leaky_relu(PF.convolution(F.concatenate(x, conv1, conv2, axis=1), growth_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv3'), alpha=0.2, inplace=True) conv4 = F.leaky_relu(PF.convolution(F.concatenate(x, conv1, conv2, conv3, axis=1), growth_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv4'), alpha=0.2, inplace=True) conv5 = PF.convolution(F.concatenate(x, conv1, conv2, conv3, conv4, axis=1), num_output_channel, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='conv5') return (conv5 * 0.2) + x
def discriminator_block(input, out_channels, scope, kernel=(4, 4), stride=(2, 2)): with nn.parameter_scope(scope): net = conv2d(input, out_channels, kernel, stride, bias=False) net = PF.batch_normalization(net, eps=0.001, axes=[3], no_scale=True) net = F.leaky_relu(net, alpha=0.2) return net
def conv_layer(conv_input, inmaps, outmaps, kernel_size, downsample=False, bias=True, act=F.leaky_relu, name_scope='Conv'): """ Conv layer for the residual block of the discriminator """ if downsample: k = [1, 3, 3, 1] out = downsample_2d(conv_input, k, factor=2, gain=1, kernel_size=kernel_size) stride = 2 pad = 0 else: stride = 1 pad = kernel_size // 2 out = conv_input init_function = weight_init_fn(shape=(outmaps, inmaps, kernel_size, kernel_size), return_init=True) scale = 1 / np.sqrt(inmaps * kernel_size**2) conv_weight = nn.parameter.get_parameter_or_create( name=f'{name_scope}/W', initializer=init_function, shape=(outmaps, inmaps, kernel_size, kernel_size)) if bias: conv_bias = nn.parameter.get_parameter_or_create( name=f'{name_scope}/b', shape=(outmaps, )) else: conv_bias = None out = F.convolution(out, conv_weight * scale, bias=conv_bias, stride=(stride, stride), pad=(pad, pad)) if act == F.leaky_relu: out = F.mul_scalar(F.leaky_relu(out, alpha=0.2, inplace=False), np.sqrt(2), inplace=False) else: out = act(out) return out
def easy_pcd(feature_p1, feature_p2, n_filt, name): """ easy 3 level pyramid cascade aligning input: features (feature_p1, feature_p2) feature size: f1 = f2 = [B, N, C, H, W] """ with nn.parameter_scope(name): # L1: level 1, original spatial size l1_fea = F.stack(*[feature_p1, feature_p2], axis=1) batch, num_frames, channels, height, width = l1_fea.shape l1_fea = l1_fea.reshape((-1, channels, height, width)) # L2: level 2, 1/2 spatial size l2_fea = F.leaky_relu(conv2d(l1_fea, n_filt, 3, 2, 1, bias=True, name='fea_l2_conv1') ) l2_fea = F.leaky_relu(conv2d(l2_fea, n_filt, 3, 1, 1, bias=True, name='fea_l2_conv2') ) # L3: level 3, 1/4 spatial size l3_fea = F.leaky_relu(conv2d(l2_fea, n_filt, 3, 2, 1, bias=True, name='fea_l3_conv1') ) l3_fea = F.leaky_relu(conv2d(l3_fea, n_filt, 3, 1, 1, bias=True, name='fea_l3_conv2') ) l1_fea = F.reshape(l1_fea, (batch, num_frames, -1, height, width), inplace=False) l2_fea = F.reshape(l2_fea, (batch, num_frames, -1, height // 2, width // 2), inplace=False) l3_fea = F.reshape(l3_fea, (batch, num_frames, -1, height // 4, width // 4), inplace=False) fea1 = [l1_fea[:, 0, :, :, :], l2_fea[:, 0, :, :, :], l3_fea[:, 0, :, :, :]] fea2 = [l1_fea[:, 1, :, :, :], l2_fea[:, 1, :, :, :], l3_fea[:, 1, :, :, :]] aligned_fea = pcd_align(fea1, fea2) fusion_fea = conv2d(aligned_fea, n_filt, 1, 1, 0, bias=True, name='fusion') return fusion_fea
def feature_extractor(input_variable, input_nc, ndf, n_layers, kw, padw, norm_layer, use_bias): w_init = I.NormalInitializer(sigma=0.02, rng=None) x = input_variable with nn.parameter_scope("feature_extractor_init_conv"): def apply_w(w): return PF.spectral_norm(w, dim=0) h = PF.convolution(x, ndf, kernel=(kw, kw), stride=(2, 2), pad=(padw, padw), w_init=w_init, apply_w=apply_w) h = F.leaky_relu(h, alpha=0.2) nf_mult = 1 nf_mult_prev = 1 for n in range(1, n_layers): nf_mult_prev = nf_mult nf_mult = min(2**n, 8) with nn.parameter_scope("feature_extractor_stage_{}".format(n)): def apply_w(w): return PF.spectral_norm(w, dim=0) h = PF.convolution(h, ndf * nf_mult, kernel=(kw, kw), stride=(2, 2), pad=(padw, padw), w_init=w_init, with_bias=use_bias, apply_w=apply_w) h = norm_layer(h) h = F.leaky_relu(h, alpha=0.2) return h
def call(self, x, y): hp = self.hp results = [] with nn.parameter_scope('layer_0'): x = F.pad(x, (0, 0, 7, 7), 'reflect') x = wn_conv(x, hp.ndf, (15,)) x = F.leaky_relu(x, 0.2, inplace=True) results.append(x) nf = hp.ndf stride = hp.downsamp_factor for i in range(1, hp.n_layers_D + 1): nf_prev = nf nf = min(nf * stride, 1024) with nn.parameter_scope(f'layer_{i}'): x = wn_conv( x, nf, (stride * 10 + 1,), stride=(stride,), pad=(stride * 5,), group=nf_prev // 4, ) x = F.leaky_relu(x, 0.2, inplace=True) results.append(x) with nn.parameter_scope(f'layer_{hp.n_layers_D + 1}'): nf = min(nf * 2, 1024) x = wn_conv(x, nf, kernel=(5,), pad=(2,)) x = F.leaky_relu(x, 0.2, inplace=True) results.append(x) with nn.parameter_scope(f'layer_{hp.n_layers_D + 2}'): x = wn_conv(x, hp.n_speakers, kernel=(3,), pad=(1,)) if y is not None: idx = F.stack( F.arange(0, hp.batch_size), y.reshape((hp.batch_size,)) ) x = F.gather_nd(x, idx) results.append(x) return results
def convblock(x, n=0, k=(4, 4), s=(2, 2), p=(1, 1), leaky=False, init_method=None): x = convolution(x, n=n, kernel=k, stride=s, pad=p, init_method=init_method) x = instance_normalization(x, fix_parameters=True) x = F.leaky_relu(x, alpha=0.2) if leaky else F.relu(x) return x
def __call__(self, x): outs = {} feats = {} with ps("discriminator/patch_gan"): # Create all scale inputs first. inputs = [x] for i in range(self.n_scales - 1): inputs.append( F.average_pooling(inputs[-1], (3, 3), (2, 2), pad=(1, 1), including_pad=False)) for i in range(self.n_scales): # Get input in reverse order of its scale (from coarse to fine) to preserve discriminator indexes. h = inputs.pop() d_name = "d_{}".format(i) feats[d_name] = {} with ps(d_name): # first layer fdim = self.base_ndf with ps("layer_0"): h = self.pad_conv(h, fdim, stride=(2, 2)) h = F.leaky_relu(h, alpha=0.2, inplace=True) feats[d_name]["layer_0"] = h # middle for j in range(1, self.n_layers - 1): fdim = min(fdim * 2, 512) layer_name = "layer_{}".format(j) with ps(layer_name): h = self.pad_conv(h, fdim, stride=(2, 2)) h = self.instance_norm_lrelu(h) feats[d_name][layer_name] = h # last 2 layers fdim = min(fdim * 2, 512) layer_name = "layer_{}".format(self.n_layers - 1) with ps(layer_name): h = self.pad_conv(h, fdim, stride=(1, 1)) h = self.instance_norm_lrelu(h) feats[d_name][layer_name] = h with nn.parameter_scope("last_layer"): h = self.pad_conv(h, 1, stride=(1, 1)) if self.use_sigmoid: h = F.sigmoid(h) outs[d_name] = h return outs, feats
def discriminator(x, image_size=128, conv_dim=64, c_dim=5, repeat_num=6, w_init=None): assert x.shape[-1] == image_size, "image_size and input spatial size must match." h = PF.convolution(x, conv_dim, kernel=(4, 4), pad=( 1, 1), stride=(2, 2), w_init=w_init, name="init_conv") h = F.leaky_relu(h, alpha=0.01) curr_dim = conv_dim for i in range(1, repeat_num): h = PF.convolution(h, curr_dim*2, kernel=(4, 4), pad=(1, 1), stride=(2, 2), w_init=w_init, name="downsample_{}".format(i)) h = F.leaky_relu(h, alpha=0.01) curr_dim = curr_dim * 2 kernel_size = int(image_size / np.power(2, repeat_num)) out_src = PF.convolution(h, 1, kernel=(3, 3), pad=(1, 1), stride=( 1, 1), with_bias=False, w_init=w_init, name="last_1x1_conv") out_cls = PF.convolution(h, c_dim, kernel=( kernel_size, kernel_size), with_bias=False, w_init=w_init, name="last_cls_conv") return out_src, out_cls