mirror of https://github.com/coqui-ai/TTS.git
85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
import torch
|
|
from torch import nn
|
|
|
|
|
|
class TimeDepthSeparableConv(nn.Module):
|
|
"""Time depth separable convolution as in https://arxiv.org/pdf/1904.02619.pdf
|
|
It shows competative results with less computation and memory footprint."""
|
|
|
|
def __init__(self, in_channels, hid_channels, out_channels, kernel_size, bias=True):
|
|
super().__init__()
|
|
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.hid_channels = hid_channels
|
|
self.kernel_size = kernel_size
|
|
|
|
self.time_conv = nn.Conv1d(
|
|
in_channels,
|
|
2 * hid_channels,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=bias,
|
|
)
|
|
self.norm1 = nn.BatchNorm1d(2 * hid_channels)
|
|
self.depth_conv = nn.Conv1d(
|
|
hid_channels,
|
|
hid_channels,
|
|
kernel_size,
|
|
stride=1,
|
|
padding=(kernel_size - 1) // 2,
|
|
groups=hid_channels,
|
|
bias=bias,
|
|
)
|
|
self.norm2 = nn.BatchNorm1d(hid_channels)
|
|
self.time_conv2 = nn.Conv1d(
|
|
hid_channels,
|
|
out_channels,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=bias,
|
|
)
|
|
self.norm3 = nn.BatchNorm1d(out_channels)
|
|
|
|
def forward(self, x):
|
|
x_res = x
|
|
x = self.time_conv(x)
|
|
x = self.norm1(x)
|
|
x = nn.functional.glu(x, dim=1)
|
|
x = self.depth_conv(x)
|
|
x = self.norm2(x)
|
|
x = x * torch.sigmoid(x)
|
|
x = self.time_conv2(x)
|
|
x = self.norm3(x)
|
|
x = x_res + x
|
|
return x
|
|
|
|
|
|
class TimeDepthSeparableConvBlock(nn.Module):
|
|
def __init__(self, in_channels, hid_channels, out_channels, num_layers, kernel_size, bias=True):
|
|
super().__init__()
|
|
assert (kernel_size - 1) % 2 == 0
|
|
assert num_layers > 1
|
|
|
|
self.layers = nn.ModuleList()
|
|
layer = TimeDepthSeparableConv(
|
|
in_channels, hid_channels, out_channels if num_layers == 1 else hid_channels, kernel_size, bias
|
|
)
|
|
self.layers.append(layer)
|
|
for idx in range(num_layers - 1):
|
|
layer = TimeDepthSeparableConv(
|
|
hid_channels,
|
|
hid_channels,
|
|
out_channels if (idx + 1) == (num_layers - 1) else hid_channels,
|
|
kernel_size,
|
|
bias,
|
|
)
|
|
self.layers.append(layer)
|
|
|
|
def forward(self, x, mask):
|
|
for layer in self.layers:
|
|
x = layer(x * mask)
|
|
return x
|