Spaces:
Running
Running
| import math | |
| import unittest | |
| import torch | |
| from models.config import VocoderModelConfig | |
| from models.vocoder.univnet import DiscriminatorP | |
| class TestDiscriminatorP(unittest.TestCase): | |
| def setUp(self): | |
| self.batch_size = 2 | |
| self.channels = 1 | |
| self.time_steps = 100 | |
| self.period = 10 | |
| self.model_config = VocoderModelConfig() | |
| self.x = torch.randn(self.batch_size, self.channels, self.time_steps) | |
| self.model = DiscriminatorP(self.period, self.model_config) | |
| def test_forward(self): | |
| fmap, output = self.model(self.x) | |
| self.assertEqual(len(fmap), len(self.model.convs) + 1) | |
| # Assert the shape of the feature maps explicitly | |
| fmap_dims = [ | |
| torch.Size([self.batch_size, 64, 4, self.period]), | |
| torch.Size([self.batch_size, 128, 2, self.period]), | |
| torch.Size([self.batch_size, 256, 1, self.period]), | |
| torch.Size([self.batch_size, 512, 1, self.period]), | |
| torch.Size([self.batch_size, 1024, 1, self.period]), | |
| torch.Size([self.batch_size, 1, 1, self.period]), | |
| ] | |
| for i in range(len(fmap)): | |
| self.assertEqual(fmap[i].shape, fmap_dims[i]) | |
| # Assert the shape of the feature maps | |
| dim_2nd = 4 | |
| for i in range(len(self.model_config.mpd.periods)): | |
| self.assertEqual(fmap[i].shape[0], self.batch_size) | |
| self.assertEqual(fmap[i].shape[1], 2 ** (i + 6)) | |
| self.assertEqual(fmap[i].shape[2], dim_2nd) | |
| dim_2nd = math.ceil(dim_2nd / self.model_config.mpd.stride) | |
| self.assertEqual(fmap[i].shape[3], self.period) | |
| self.assertEqual(output.shape, (self.batch_size, self.period)) | |
| def test_forward_with_padding(self): | |
| fmap, output = self.model(self.x) | |
| self.assertEqual(len(fmap), len(self.model.convs) + 1) | |
| # Assert the shape of the feature maps explicitly | |
| fmap_dims = [ | |
| torch.Size([self.batch_size, 64, 4, self.period]), | |
| torch.Size([self.batch_size, 128, 2, self.period]), | |
| torch.Size([self.batch_size, 256, 1, self.period]), | |
| torch.Size([self.batch_size, 512, 1, self.period]), | |
| torch.Size([self.batch_size, 1024, 1, self.period]), | |
| torch.Size([self.batch_size, 1, 1, self.period]), | |
| ] | |
| for i in range(len(fmap)): | |
| self.assertEqual(fmap[i].shape, fmap_dims[i]) | |
| # Assert the shape of the feature maps | |
| dim_2nd = 4 | |
| for i in range(len(self.model_config.mpd.periods)): | |
| self.assertEqual(fmap[i].shape[0], self.batch_size) | |
| self.assertEqual(fmap[i].shape[1], 2 ** (i + 6)) | |
| self.assertEqual(fmap[i].shape[2], dim_2nd) | |
| dim_2nd = math.ceil(dim_2nd / self.model_config.mpd.stride) | |
| self.assertEqual(fmap[i].shape[3], self.period) | |
| self.assertEqual(output.shape, (self.batch_size, self.period)) | |
| def test_forward_with_different_period(self): | |
| model = DiscriminatorP(self.period, self.model_config) | |
| x = torch.randn(self.batch_size, self.channels, self.time_steps - 1) | |
| model.period = 5 | |
| fmap, output = model(x) | |
| self.assertEqual(len(fmap), len(model.convs) + 1) | |
| # Assert the shape of the feature maps explicitly | |
| fmap_dims = [ | |
| torch.Size([self.batch_size, 64, 7, model.period]), | |
| torch.Size([self.batch_size, 128, 3, model.period]), | |
| torch.Size([self.batch_size, 256, 1, model.period]), | |
| torch.Size([self.batch_size, 512, 1, model.period]), | |
| torch.Size([self.batch_size, 1024, 1, model.period]), | |
| torch.Size([self.batch_size, 1, 1, model.period]), | |
| ] | |
| for i in range(len(fmap)): | |
| self.assertEqual(fmap[i].shape, fmap_dims[i]) | |
| # Assert the shape of the feature maps | |
| dim_2nd = 7 | |
| for i in range(len(self.model_config.mpd.periods)): | |
| self.assertEqual(fmap[i].shape[0], self.batch_size) | |
| self.assertEqual(fmap[i].shape[1], 2 ** (i + 6)) | |
| self.assertEqual(fmap[i].shape[2], dim_2nd) | |
| dim_2nd = math.ceil(dim_2nd / self.model_config.mpd.stride) | |
| self.assertEqual(fmap[i].shape[3], model.period) | |
| self.assertEqual(output.shape, (self.batch_size, model.period)) | |
| if __name__ == "__main__": | |
| unittest.main() | |