mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-27 08:37:40 +00:00
first commit; melgan support (perhaps bloated)
I copied this from my old and messy `tacotron-melgan` repo, where I had previously experimented with running the two models in a pipeline. I'll be moving over all the JIT pieces here and cleaning up when I can.
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .discriminator import Discriminator
|
||||
from .identity import Identity
|
||||
|
||||
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
def __init__(self):
|
||||
super(MultiScaleDiscriminator, self).__init__()
|
||||
|
||||
self.discriminators = nn.ModuleList(
|
||||
[Discriminator() for _ in range(3)]
|
||||
)
|
||||
|
||||
self.pooling = nn.ModuleList(
|
||||
[Identity()] +
|
||||
[nn.AvgPool1d(kernel_size=4, stride=2, padding=1, count_include_pad=False) for _ in range(1, 3)]
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
ret = list()
|
||||
|
||||
for pool, disc in zip(self.pooling, self.discriminators):
|
||||
x = pool(x)
|
||||
ret.append(disc(x))
|
||||
|
||||
return ret # [(feat, score), (feat, score), (feat, score)]
|
||||
Reference in New Issue
Block a user