mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-18 07:33:20 +00:00
Brandon Thomas
245 lines
8.2 KiB
Python
245 lines
8.2 KiB
Python
import copy
|
|
import math
|
|
import numpy as np
|
|
import scipy
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
import commons
|
|
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, channels, eps=1e-4):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.eps = eps
|
|
|
|
self.gamma = nn.Parameter(torch.ones(channels))
|
|
self.beta = nn.Parameter(torch.zeros(channels))
|
|
|
|
def forward(self, x):
|
|
n_dims = len(x.shape)
|
|
mean = torch.mean(x, 1, keepdim=True)
|
|
variance = torch.mean((x -mean)**2, 1, keepdim=True)
|
|
|
|
x = (x - mean) * torch.rsqrt(variance + self.eps)
|
|
|
|
# TODO(bt): Let's make this a static view.
|
|
# This isn't correct, but maybe we can put branch cases in for every possible length.
|
|
# It depends on the value of n_dims
|
|
shape = [1, -1] + [1] * (n_dims - 2)
|
|
shape = [1, -1, 1]
|
|
#x = x * self.gamma.view(*shape) + self.beta.view(*shape)
|
|
x = x * self.gamma.view(shape[0], shape[1], shape[2]) + self.beta.view(shape[0], shape[1], shape[2])
|
|
return x
|
|
|
|
|
|
class ConvReluNorm(nn.Module):
|
|
def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.out_channels = out_channels
|
|
self.kernel_size = kernel_size
|
|
self.n_layers = n_layers
|
|
self.p_dropout = p_dropout
|
|
assert n_layers > 1, "Number of layers should be larger than 0."
|
|
|
|
self.conv_layers = nn.ModuleList()
|
|
self.norm_layers = nn.ModuleList()
|
|
self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
|
|
self.norm_layers.append(LayerNorm(hidden_channels))
|
|
self.relu_drop = nn.Sequential(
|
|
nn.ReLU(),
|
|
nn.Dropout(p_dropout))
|
|
for _ in range(n_layers-1):
|
|
self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
|
|
self.norm_layers.append(LayerNorm(hidden_channels))
|
|
self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
|
|
self.proj.weight.data.zero_()
|
|
self.proj.bias.data.zero_()
|
|
|
|
def forward(self, x, x_mask):
|
|
x_org = x
|
|
# NB(bt): Fixed for torchscript
|
|
for conv_layer, norm_layer in zip(self.conv_layers, self.norm_layers):
|
|
x = conv_layer(x * x_mask)
|
|
x = norm_layer(x)
|
|
x = self.relu_drop(x)
|
|
x = x_org + self.proj(x)
|
|
return x * x_mask
|
|
|
|
|
|
class WN(torch.nn.Module):
|
|
def __init__(self, in_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
|
|
super(WN, self).__init__()
|
|
assert(kernel_size % 2 == 1)
|
|
assert(hidden_channels % 2 == 0)
|
|
self.in_channels = in_channels
|
|
self.hidden_channels =hidden_channels
|
|
self.kernel_size = kernel_size,
|
|
self.dilation_rate = dilation_rate
|
|
self.n_layers = n_layers
|
|
self.gin_channels = gin_channels
|
|
self.p_dropout = p_dropout
|
|
|
|
self.in_layers = torch.nn.ModuleList()
|
|
self.res_skip_layers = torch.nn.ModuleList()
|
|
self.drop = nn.Dropout(p_dropout)
|
|
|
|
if gin_channels != 0:
|
|
cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
|
|
self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
|
|
|
|
for i in range(n_layers):
|
|
dilation = dilation_rate ** i
|
|
padding = int((kernel_size * dilation - dilation) / 2)
|
|
in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
|
|
dilation=dilation, padding=padding)
|
|
in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
|
|
self.in_layers.append(in_layer)
|
|
|
|
# last one is not necessary
|
|
if i < n_layers - 1:
|
|
res_skip_channels = 2 * hidden_channels
|
|
else:
|
|
res_skip_channels = hidden_channels
|
|
|
|
res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
|
|
res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
|
|
self.res_skip_layers.append(res_skip_layer)
|
|
|
|
def forward(self, x, x_mask=None, g=None, **kwargs):
|
|
output = torch.zeros_like(x)
|
|
n_channels_tensor = torch.IntTensor([self.hidden_channels])
|
|
|
|
if g is not None:
|
|
g = self.cond_layer(g)
|
|
|
|
for i in range(self.n_layers):
|
|
x_in = self.in_layers[i](x)
|
|
x_in = self.drop(x_in)
|
|
if g is not None:
|
|
cond_offset = i * 2 * self.hidden_channels
|
|
g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
|
|
else:
|
|
g_l = torch.zeros_like(x_in)
|
|
|
|
acts = commons.fused_add_tanh_sigmoid_multiply(
|
|
x_in,
|
|
g_l,
|
|
n_channels_tensor)
|
|
|
|
res_skip_acts = self.res_skip_layers[i](acts)
|
|
if i < self.n_layers - 1:
|
|
x = (x + res_skip_acts[:,:self.hidden_channels,:]) * x_mask
|
|
output = output + res_skip_acts[:,self.hidden_channels:,:]
|
|
else:
|
|
output = output + res_skip_acts
|
|
return output * x_mask
|
|
|
|
def remove_weight_norm(self):
|
|
if self.gin_channels != 0:
|
|
torch.nn.utils.remove_weight_norm(self.cond_layer)
|
|
for l in self.in_layers:
|
|
torch.nn.utils.remove_weight_norm(l)
|
|
for l in self.res_skip_layers:
|
|
torch.nn.utils.remove_weight_norm(l)
|
|
|
|
|
|
class ActNorm(nn.Module):
|
|
def __init__(self, channels, ddi=False, **kwargs):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.initialized = not ddi
|
|
|
|
self.logs = nn.Parameter(torch.zeros(1, channels, 1))
|
|
self.bias = nn.Parameter(torch.zeros(1, channels, 1))
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, **kwargs):
|
|
if x_mask is None:
|
|
x_mask = torch.ones(x.size(0), 1, x.size(2)).to(device=x.device, dtype=x.dtype)
|
|
x_len = torch.sum(x_mask, [1, 2])
|
|
if not self.initialized:
|
|
self.initialize(x, x_mask)
|
|
self.initialized = True
|
|
|
|
if reverse:
|
|
z = (x - self.bias) * torch.exp(-self.logs) * x_mask
|
|
logdet = None
|
|
else:
|
|
z = (self.bias + torch.exp(self.logs) * x) * x_mask
|
|
logdet = torch.sum(self.logs) * x_len # [b]
|
|
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
pass
|
|
|
|
def set_ddi(self, ddi):
|
|
self.initialized = not ddi
|
|
|
|
def initialize(self, x, x_mask):
|
|
with torch.no_grad():
|
|
denom = torch.sum(x_mask, [0, 2])
|
|
m = torch.sum(x * x_mask, [0, 2]) / denom
|
|
m_sq = torch.sum(x * x * x_mask, [0, 2]) / denom
|
|
v = m_sq - (m ** 2)
|
|
logs = 0.5 * torch.log(torch.clamp_min(v, 1e-6))
|
|
|
|
bias_init = (-m * torch.exp(-logs)).view(*self.bias.shape).to(dtype=self.bias.dtype)
|
|
logs_init = (-logs).view(*self.logs.shape).to(dtype=self.logs.dtype)
|
|
|
|
self.bias.data.copy_(bias_init)
|
|
self.logs.data.copy_(logs_init)
|
|
|
|
|
|
class InvConvNear(nn.Module):
|
|
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
|
|
super().__init__()
|
|
assert(n_split % 2 == 0)
|
|
self.channels = channels
|
|
self.n_split = n_split
|
|
self.no_jacobian = no_jacobian
|
|
|
|
w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
|
|
if torch.det(w_init) < 0:
|
|
w_init[:,0] = -1 * w_init[:,0]
|
|
self.weight = nn.Parameter(w_init)
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, **kwargs):
|
|
b, c, t = x.size()
|
|
assert(c % self.n_split == 0)
|
|
if x_mask is None:
|
|
x_mask = 1
|
|
x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t
|
|
else:
|
|
x_len = torch.sum(x_mask, [1, 2])
|
|
|
|
x = x.view(b, 2, c // self.n_split, self.n_split // 2, t)
|
|
x = x.permute(0, 1, 3, 2, 4).contiguous().view(b, self.n_split, c // self.n_split, t)
|
|
|
|
if reverse:
|
|
if hasattr(self, "weight_inv"):
|
|
weight = self.weight_inv
|
|
else:
|
|
weight = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
|
|
logdet = None
|
|
else:
|
|
weight = self.weight
|
|
if self.no_jacobian:
|
|
logdet = 0
|
|
else:
|
|
logdet = torch.logdet(self.weight) * (c / self.n_split) * x_len # [b]
|
|
|
|
weight = weight.view(self.n_split, self.n_split, 1, 1)
|
|
z = F.conv2d(x, weight)
|
|
|
|
z = z.view(b, 2, self.n_split // 2, c // self.n_split, t)
|
|
z = z.permute(0, 1, 3, 2, 4).contiguous().view(b, c, t) * x_mask
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
self.weight_inv = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
|