Spaces:
Sleeping
Sleeping
File size: 2,041 Bytes
14d1720 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import torch
import torch.nn as nn
from utils.utils import weights_init
from model.discriminator import JCU_Discriminator, Discriminator
class MultiScaleDiscriminator(nn.Module):
def __init__(self, num_D = 3, ndf = 16, n_layers = 3, downsampling_factor = 4, disc_out = 512):
super().__init__()
self.model = nn.ModuleDict()
for i in range(num_D):
self.model[f"disc_{i}"] = Discriminator(
ndf, n_layers, downsampling_factor, disc_out
)
self.downsample = nn.AvgPool1d(downsampling_factor, stride=2, padding=1, count_include_pad=False)
self.apply(weights_init)
def forward(self, x):
results = []
for key, disc in self.model.items():
results.append(disc(x))
x = self.downsample(x)
return results
class MultiScaleDiscriminatorJCU(nn.Module):
def __init__(self, num_D = 3, downsampling_factor = 4):
super(MultiScaleDiscriminator, self).__init__()
self.model = nn.ModuleDict()
for i in range(num_D):
self.model[f"disc_{i}"] = JCU_Discriminator()
self.downsample = nn.AvgPool1d(downsampling_factor, stride=2, padding=1, count_include_pad=False)
def forward(self, x, mel):
results = []
for key, disc in self.model.items():
results.append(disc(x, mel)) # [[uncond, cond], [uncond, cond], [uncond, cond]]
x = self.downsample(x)
mel = self.downsample(mel)
return results # [D01, D02, D03]
if __name__ == '__main__':
model = MultiScaleDiscriminator()
x = torch.randn(3, 1, 22050)
print(x.shape)
print(model)
scores = model(x)
for (features, score) in scores:
print("Length of features : ", len(features))
print("Length of score : ", len(score))
for feat in features:
print(feat.shape)
print(score.shape)
pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(pytorch_total_params) |