173 lines
8.5 KiB
Python
173 lines
8.5 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from ..commons.layers import LayerNorm, Embedding
|
|
from ..commons.conv import ConvBlocks, ResidualBlock, get_norm_builder, get_act_builder
|
|
|
|
class UnetDown(nn.Module):
|
|
def __init__(self, hidden_size, n_layers, kernel_size, down_rates, channel_multiples=None, dropout=0.0,
|
|
is_BTC=True, constant_channels=False):
|
|
super(UnetDown, self).__init__()
|
|
assert n_layers == len(down_rates) # downs, down sample rate
|
|
down_rates = [int(i) for i in down_rates]
|
|
self.n_layers = n_layers
|
|
self.hidden_size = hidden_size
|
|
self.is_BTC = is_BTC
|
|
channel_multiples = channel_multiples if channel_multiples is not None else down_rates
|
|
self.layers = nn.ModuleList()
|
|
self.downs = nn.ModuleList()
|
|
in_channels = hidden_size
|
|
for i in range(self.n_layers):
|
|
out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels
|
|
self.layers.append(nn.Sequential(
|
|
ResidualBlock(in_channels, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout,
|
|
c_multiple=1, ln_eps=1e-5, act_type='leakyrelu'),
|
|
nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2),
|
|
ResidualBlock(out_channels, kernel_size, dilation=1, n=1, norm_type='ln',
|
|
dropout=dropout, c_multiple=1, ln_eps=1e-5, act_type='leakyrelu')
|
|
))
|
|
self.downs.append(nn.Sequential(
|
|
nn.AvgPool1d(down_rates[i])
|
|
))
|
|
in_channels = out_channels
|
|
self.last_norm = get_norm_builder('ln', out_channels)()
|
|
self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size,
|
|
padding=kernel_size // 2)
|
|
|
|
def forward(self, x, **kwargs):
|
|
# x [B, T, C]
|
|
if self.is_BTC:
|
|
x = x.transpose(1, 2) # [B, C, T]
|
|
skip_xs = []
|
|
for i in range(self.n_layers):
|
|
skip_x = self.layers[i](x)
|
|
x = self.downs[i](skip_x)
|
|
if self.is_BTC:
|
|
skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C]
|
|
else:
|
|
skip_xs.append(skip_x)
|
|
x = self.post_net(self.last_norm(x))
|
|
if self.is_BTC:
|
|
x = x.transpose(1, 2)
|
|
return x, skip_xs
|
|
|
|
class UnetMid(nn.Module):
|
|
def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None,
|
|
dropout=0.0, is_BTC=True, net=None):
|
|
super(UnetMid, self).__init__()
|
|
in_dims = in_dims if in_dims is not None else hidden_size
|
|
out_dims = out_dims if out_dims is not None else hidden_size
|
|
self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2)
|
|
self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2)
|
|
self.is_BTC = is_BTC
|
|
if net is not None:
|
|
self.net = net
|
|
else:
|
|
self.net = ConvBlocks(hidden_size, out_dims=hidden_size, dilations=None, kernel_size=kernel_size,
|
|
layers_in_block=2, c_multiple=2, dropout=dropout, num_layers=n_layers,
|
|
post_net_kernel=3, act_type='leakyrelu', is_BTC=is_BTC)
|
|
|
|
def forward(self, x, cond=None, **kwargs):
|
|
# x [B, T, C]
|
|
if self.is_BTC:
|
|
x = self.pre(x.transpose(1, 2)).transpose(1, 2)
|
|
else:
|
|
x = self.pre(x)
|
|
if cond is None:
|
|
cond = 0
|
|
x = self.net(x + cond)
|
|
if self.is_BTC:
|
|
x = self.post(x.transpose(1, 2)).transpose(1, 2)
|
|
else:
|
|
x = self.post(x)
|
|
return x
|
|
|
|
class UnetUp(nn.Module):
|
|
def __init__(self, hidden_size, n_layers, kernel_size, up_rates, channel_multiples=None, dropout=0.0,
|
|
is_BTC=True, constant_channels=False, use_skip_layer=False, skip_scale=1.0):
|
|
super(UnetUp, self).__init__()
|
|
assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle
|
|
up_rates = [int(i) for i in up_rates]
|
|
self.n_layers = n_layers
|
|
self.hidden_size = hidden_size
|
|
self.is_BTC = is_BTC
|
|
self.skip_scale = skip_scale
|
|
channel_multiples = channel_multiples if channel_multiples is not None else up_rates
|
|
# in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size
|
|
self.in_channels_lst = (np.cumprod([1] + channel_multiples) * hidden_size).astype(int) if not constant_channels \
|
|
else [hidden_size for _ in range(self.n_layers + 1)]
|
|
in_channels = self.in_channels_lst[-1]
|
|
self.ups = nn.ModuleList()
|
|
self.skip_layers = nn.ModuleList()
|
|
self.layers = nn.ModuleList()
|
|
for i in range(self.n_layers-1, -1, -1):
|
|
out_channels = self.in_channels_lst[i] if not constant_channels else in_channels
|
|
self.ups.append(nn.Sequential(
|
|
nn.ConvTranspose1d(in_channels, in_channels, kernel_size=kernel_size, stride=up_rates[i],
|
|
padding=kernel_size//2, output_padding=up_rates[i]-1),
|
|
get_norm_builder('ln', in_channels)(),
|
|
get_act_builder('leakyrelu')()
|
|
))
|
|
self.layers.append(nn.Sequential(
|
|
# ResidualBlock(in_channels*2, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout,
|
|
# c_multiple=1, ln_eps=1e-5, act_type='leakyrelu'),
|
|
nn.Conv1d(in_channels*2, out_channels, kernel_size, padding=(kernel_size - 1) // 2),
|
|
ResidualBlock(out_channels, kernel_size, dilation=1, n=1, norm_type='ln',
|
|
dropout=dropout, c_multiple=1, ln_eps=1e-5, act_type='leakyrelu')
|
|
))
|
|
if use_skip_layer:
|
|
self.skip_layers.append(
|
|
ResidualBlock(in_channels, kernel_size, dilation=1, n=1, norm_type='ln', dropout=dropout,
|
|
c_multiple=1, ln_eps=1e-5, act_type='leakyrelu')
|
|
)
|
|
else:
|
|
self.skip_layers.append(nn.Identity())
|
|
|
|
in_channels = out_channels
|
|
self.out_channels = out_channels
|
|
self.last_norm = get_norm_builder('ln', out_channels)()
|
|
self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size,
|
|
padding=kernel_size // 2)
|
|
|
|
def forward(self, x, skips, **kwargs):
|
|
# x [B, T, C]
|
|
if self.is_BTC:
|
|
x = x.transpose(1, 2) # [B, C, T]
|
|
for i in range(self.n_layers):
|
|
x = self.ups[i](x)
|
|
skip_x = skips[self.n_layers - i - 1] if not self.is_BTC \
|
|
else skips[self.n_layers - i - 1].transpose(1, 2) # [B, T, C] -> [B, C, T]
|
|
skip_x = self.skip_layers[i](skip_x) * self.skip_scale
|
|
x = torch.cat((x, skip_x), dim=1) # [B, C, T]
|
|
x = self.layers[i](x)
|
|
x = self.post_net(self.last_norm(x))
|
|
if self.is_BTC:
|
|
x = x.transpose(1, 2)
|
|
return x
|
|
|
|
class Unet(nn.Module):
|
|
def __init__(self, hidden_size, down_layers, up_layers, kernel_size,
|
|
updown_rates, mid_layers=None, channel_multiples=None, dropout=0.0,
|
|
is_BTC=True, constant_channels=False, mid_net=None, use_skip_layer=False, skip_scale=1.0):
|
|
super(Unet, self).__init__()
|
|
assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}"
|
|
if channel_multiples is not None:
|
|
assert len(channel_multiples) == len(updown_rates)
|
|
else:
|
|
channel_multiples = updown_rates
|
|
self.down = UnetDown(hidden_size, down_layers, kernel_size, updown_rates,
|
|
channel_multiples, dropout, is_BTC, constant_channels)
|
|
down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size
|
|
self.mid = UnetMid(hidden_size, kernel_size, mid_layers,
|
|
in_dims=down_out_dims, out_dims=down_out_dims, dropout=dropout, is_BTC=is_BTC, net=mid_net)
|
|
self.up = UnetUp(hidden_size, up_layers, kernel_size, updown_rates,
|
|
channel_multiples, dropout, is_BTC, constant_channels, use_skip_layer, skip_scale)
|
|
|
|
def forward(self, x, mid_cond=None, **kwargs):
|
|
x, skips = self.down(x)
|
|
x = self.mid(x, mid_cond)
|
|
x = self.up(x, skips)
|
|
return x
|