Update gan_model.py

gan_model.py

@@ -1,5 +1,3 @@
-import logging
-
 import torch
 from torch import nn
 from torch.nn import init
@@ -216,105 +214,3 @@ class Generator(nn.Module):
         out = self.colorize(out) # [bsz, 1, 80, 344]
 
         return out
-
-
-class Discriminator(nn.Module):
-    def __init__(self, in_channel=1, channel=32, num_classes=7, embedding_dim=128):
-        super().__init__()
-        self.num_classes = num_classes
-
-        def conv(in_channel, out_channel, downsample=True):
-            return ConvBlock(in_channel, out_channel,
-                             bn=False,
-                             upsample=False, downsample=downsample)
-
-        gain = 2 ** 0.5
-
-        self.pre_conv = nn.Sequential(spectral_init(nn.Conv2d(in_channel, channel, 3,
-                                                              padding=1),
-                                                    gain=gain),
-                                      nn.ReLU(),
-                                      spectral_init(nn.Conv2d(channel, channel, 3,
-                                                              padding=1),
-                                                    gain=gain),
-                                      nn.AvgPool2d(2))
-        self.pre_skip = spectral_init(nn.Conv2d(in_channel, channel, 1))
-
-        self.conv1 = conv(channel, channel * 2)
-        self.conv2 = conv(channel * 2, channel * 2, downsample=False)
-        self.attention = SelfAttention(channel * 2)
-        self.conv3 = conv(channel * 2, channel * 4)
-        self.conv4 = conv(channel * 4, channel * 4)
-
-        self.linear = spectral_init(nn.Linear(channel * 4, 1))
-
-        self.projection = nn.Sequential(
-            spectral_init(nn.Linear(channel * 4, channel * 4)),
-            nn.ReLU(),
-            spectral_init(nn.Linear(channel * 4, channel * 4))
-        )
-
-        self.embedding = spectral_norm(nn.Embedding(num_embeddings=num_classes, embedding_dim=channel * 4))
-
-    def forward(self, input, label):
-        out = self.pre_conv(input) # [bsz, 32, 40, 172]
-        out = out + self.pre_skip(F.avg_pool2d(input, 2)) # [bsz, 32, 40, 172]
-
-        out = self.conv1(out) # [bsz, 64, 20, 86]
-        out = self.conv2(out) # [bsz, 64, 20, 86]
-        out, attention_map = self.attention(out) # [bsz, 64, 20, 86]
-        out = self.conv3(out) # [bsz, 128, 10, 43]
-        out = self.conv4(out) # [bsz, 128, 5, 21]
-
-        out = F.relu(out)
-        out = out.view(out.size(0), out.size(1), -1) # [bsz, 128, 105]
-        out = out.sum(2) # [bsz, 128]
-        adv_output = self.linear(out).squeeze(1) # [bsz, 1]
-
-        condition = self.embedding(label) # [bsz, 128]
-        prod = (out * condition).sum(1) # [bsz, 1]
-        adv_output += prod
-
-        contrastive_feature = self.projection(out) # [bsz, 128]
-
-        return adv_output, contrastive_feature, condition
-
-def count_parameters(module):
-    num_params = sum(p.numel() for p in module.parameters())
-    return num_params
-
-if __name__ == '__main__':
-    from HiFiGanWrapper import HiFiGanWrapper
-    import numpy as np
-
-    generator = Generator().eval()
-    num_params = count_parameters(generator)
-    print(f"Number of generator parameters: {num_params / 1000000:.2f} M")
-    print()
-
-    discriminator = Discriminator().eval()
-    num_params = count_parameters(discriminator)
-    print(f"Number of discriminator parameters: {num_params / 1000000:.2f} M")
-    print()
-
-    vocoder = HiFiGanWrapper(ckpt_path='./pretrained_checkpoints')
-    num_params = count_parameters(vocoder.generator)
-    print(f"Number of vocoder parameters: {num_params / 1000000:.2f} M")
-    print()
-
-    image = torch.randn(4, 1, 80, 344)
-    labels = torch.LongTensor([0, 0, 1, 2])
-
-    out, contrastive_feature, proxy = discriminator(image, labels)
-    print('discriminator :', out.shape)
-    print('contrastive_feature :', contrastive_feature.shape)
-    print('proxy :', proxy.shape)
-    print()
-
-    out = generator(labels)
-    print('generator :', out.shape)
-    print()
-
-    fake_sound = vocoder.generate_audio(out[0])
-    fake_sound = np.concatenate((fake_sound, fake_sound[-136:]), axis=0)
-    print('generated sound :', fake_sound.shape)