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_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_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]))