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
