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')))