mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-19 07:42:59 +00:00
Brandon Thomas
72dd12b460
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.
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
|
|
class ResStack(nn.Module):
|
|
def __init__(self, channel):
|
|
super(ResStack, self).__init__()
|
|
|
|
block_list = [
|
|
nn.Sequential(
|
|
nn.LeakyReLU(0.2),
|
|
nn.ReflectionPad1d(3**i),
|
|
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=3**i)),
|
|
nn.LeakyReLU(0.2),
|
|
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
|
|
)
|
|
for i in range(3)
|
|
]
|
|
|
|
self.block0 = block_list[0]
|
|
self.block1 = block_list[1]
|
|
self.block2 = block_list[2]
|
|
|
|
shortcut_list = [
|
|
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
|
|
for i in range(3)
|
|
]
|
|
|
|
self.shortcut0 = shortcut_list[0]
|
|
self.shortcut1 = shortcut_list[1]
|
|
self.shortcut2 = shortcut_list[2]
|
|
|
|
def forward(self, x):
|
|
x = self.shortcut0(x) + self.block0(x)
|
|
x = self.shortcut1(x) + self.block1(x)
|
|
x = self.shortcut2(x) + self.block2(x)
|
|
return x
|
|
|