| import torch | |
| import torch.nn as nn | |
| from omegaconf import OmegaConf | |
| from .msd import ScaleDiscriminator | |
| from .mpd import MultiPeriodDiscriminator | |
| from .mrd import MultiResolutionDiscriminator | |
| class Discriminator(nn.Module): | |
| def __init__(self, hp): | |
| super(Discriminator, self).__init__() | |
| self.MRD = MultiResolutionDiscriminator(hp) | |
| self.MPD = MultiPeriodDiscriminator(hp) | |
| self.MSD = ScaleDiscriminator() | |
| def forward(self, x): | |
| r = self.MRD(x) | |
| p = self.MPD(x) | |
| s = self.MSD(x) | |
| return r + p + s | |
| if __name__ == '__main__': | |
| hp = OmegaConf.load('../config/base.yaml') | |
| model = Discriminator(hp) | |
| x = torch.randn(3, 1, 16384) | |
| print(x.shape) | |
| output = model(x) | |
| for features, score in output: | |
| 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) | |