def build_gen_out_sequence(self):
return Sequential(
Reshape([[0], [1], [2], 1]),
weight_scale(nn.ConvTranspose2d(self.n_filters, 1,
(1, self.n_channels)),
gain=calculate_gain('linear')),
Reshape([[0], [2], [3]]), Permute([0, 2, 1]))
def build_interpolate(scale_factor: float, mode: str):
if mode in ['nearest', 'linear', 'area']:
return Interpolate(scale_factor=scale_factor, mode=mode)
if mode == 'cubic':
return Sequential(
Reshape([[0], [1], [2], 1]),
Interpolate(scale_factor=(scale_factor, 1), mode='bicubic'),
Reshape([[0], [1], [2]]))
def build_disc_conv_sequence(self, i_stage: int):
return Sequential(
weight_scale(create_multiconv_for_stage(self.n_filters, i_stage),
gain=calculate_gain('leaky_relu')), nn.LeakyReLU(0.2),
weight_scale(nn.Conv1d(self.n_filters,
self.n_filters,
kernel_size=1),
gain=calculate_gain('leaky_relu')), nn.LeakyReLU(0.2),
self.build_disc_downsample_sequence(),
weight_scale(EmbeddedClassStyle(self.n_classes, self.n_filters),
gain=calculate_gain('leaky_relu')), nn.LeakyReLU(0.2))
def build_gen_upsample_sequence(self) -> nn.Module:
if self.upsampling in ['nearest', 'linear', 'area', 'cubic']:
return build_interpolate(2, self.upsampling)
if self.upsampling == 'conv':
return Sequential(
weight_scale(nn.ConvTranspose1d(self.n_filters,
self.n_filters,
4,
stride=2,
padding=1),
gain=calculate_gain('leaky_relu')),
nn.LeakyReLU(0.2))
def build_disc_downsample_sequence(self) -> nn.Module:
if self.downsampling in ['nearest', 'linear', 'area', 'cubic']:
return build_interpolate(0.5, self.downsampling)
if self.downsampling == 'conv':
return Sequential(
nn.ReflectionPad1d(1),
weight_scale(nn.Conv1d(self.n_filters,
self.n_filters,
4,
stride=2),
gain=calculate_gain('leaky_relu')),
nn.LeakyReLU(0.2))
def build_discriminator(self) -> ProgressiveDiscriminator:
blocks = []
for i in range(self.n_stages - 1):
block = ProgressiveDiscriminatorBlock(
self.build_disc_conv_sequence(self.n_stages - 1 - i),
self.build_disc_in_sequence(), self.build_disc_fade_sequence())
blocks.append(block)
last_block = ProgressiveDiscriminatorBlock(
Sequential(
self.build_disc_conv_sequence(0),
Reshape([[0], self.n_filters * self.n_time_last_layer]),
weight_scale(nn.Linear(self.n_filters * self.n_time_last_layer,
1),
gain=calculate_gain('linear'))),
self.build_disc_in_sequence(), None)
blocks.append(last_block)
return ProgressiveDiscriminator(self.n_time, self.n_channels,
self.n_classes, blocks)
def build_generator(self) -> ProgressiveGenerator:
blocks = []
first_block = ProgressiveGeneratorBlock(
Sequential(
weight_scale(nn.Linear(self.n_latent, self.n_filters *
self.n_time_last_layer),
gain=calculate_gain('leaky_relu')),
Reshape([[0], self.n_filters, -1]), nn.LeakyReLU(0.2),
PixelNorm(), self.build_gen_conv_sequence(0)),
self.build_gen_out_sequence(), self.build_gen_fade_sequence())
blocks.append(first_block)
for i in range(1, 6):
block = ProgressiveGeneratorBlock(self.build_gen_conv_sequence(i),
self.build_gen_out_sequence(),
self.build_gen_fade_sequence())
blocks.append(block)
return ProgressiveGenerator(self.n_time, self.n_channels,
self.n_classes, self.n_latent, blocks)
def build_disc_in_sequence(self):
return Sequential(
Permute([0, 2, 1]), Reshape([[0], 1, [1], [2]]),
weight_scale(nn.Conv2d(1, self.n_filters, (1, self.n_channels)),
gain=calculate_gain('leaky_relu')),
Reshape([[0], [1], [2]]), nn.LeakyReLU(0.2))
def build_disc_in_sequence(self):
return Sequential(
weight_scale(nn.Conv1d(self.n_channels, self.n_filters, 1),
gain=calculate_gain('leaky_relu')), nn.LeakyReLU(0.2))
def build_gen_out_sequence(self):
return Sequential(
weight_scale(nn.Conv1d(self.n_filters, self.n_channels, 1),
gain=calculate_gain('linear')),
Reshape([[0], [1], [2], 1]), PixelShuffle2d((1, self.n_channels)),
Reshape([[0], [2], [3]]), Permute([0, 2, 1]))
def build_gen_out_sequence(self):
return Sequential(
weight_scale(nn.Conv1d(self.n_filters, self.n_channels, 1),
gain=calculate_gain('linear')))
