|
|
|
|
|
import math |
|
|
|
import torch |
|
import torch.nn as nn |
|
from torch import nn |
|
from torch.nn.utils.parametrizations import weight_norm |
|
|
|
from .alias_free_torch import DownSample1d, UpSample1d |
|
|
|
|
|
class SnakeBeta(nn.Module): |
|
""" |
|
A modified Snake function which uses separate parameters for the magnitude of the periodic components |
|
Shape: |
|
- Input: (B, C, T) |
|
- Output: (B, C, T), same shape as the input |
|
Parameters: |
|
- alpha - trainable parameter that controls frequency |
|
- beta - trainable parameter that controls magnitude |
|
References: |
|
- This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: |
|
https://arxiv.org/abs/2006.08195 |
|
Examples: |
|
>>> a1 = snakebeta(256) |
|
>>> x = torch.randn(256) |
|
>>> x = a1(x) |
|
""" |
|
|
|
def __init__(self, in_features, alpha=1.0, clamp=(1e-2, 50)): |
|
""" |
|
Initialization. |
|
INPUT: |
|
- in_features: shape of the input |
|
- alpha - trainable parameter that controls frequency |
|
- beta - trainable parameter that controls magnitude |
|
alpha is initialized to 1 by default, higher values = higher-frequency. |
|
beta is initialized to 1 by default, higher values = higher-magnitude. |
|
alpha will be trained along with the rest of your model. |
|
""" |
|
super().__init__() |
|
self.in_features = in_features |
|
self.log_alpha = nn.Parameter(torch.zeros(in_features) + math.log(alpha)) |
|
self.log_beta = nn.Parameter(torch.zeros(in_features) + math.log(alpha)) |
|
self.clamp = clamp |
|
|
|
def forward(self, x): |
|
""" |
|
Forward pass of the function. |
|
Applies the function to the input elementwise. |
|
SnakeBeta ∶= x + 1/b * sin^2 (xa) |
|
""" |
|
alpha = self.log_alpha.exp().clamp(*self.clamp) |
|
alpha = alpha[None, :, None] |
|
|
|
beta = self.log_beta.exp().clamp(*self.clamp) |
|
beta = beta[None, :, None] |
|
|
|
x = x + (1.0 / beta) * (x * alpha).sin().pow(2) |
|
|
|
return x |
|
|
|
|
|
class UpActDown(nn.Module): |
|
def __init__( |
|
self, |
|
act, |
|
up_ratio: int = 2, |
|
down_ratio: int = 2, |
|
up_kernel_size: int = 12, |
|
down_kernel_size: int = 12, |
|
): |
|
super().__init__() |
|
self.up_ratio = up_ratio |
|
self.down_ratio = down_ratio |
|
self.act = act |
|
self.upsample = UpSample1d(up_ratio, up_kernel_size) |
|
self.downsample = DownSample1d(down_ratio, down_kernel_size) |
|
|
|
def forward(self, x): |
|
|
|
x = self.upsample(x) |
|
x = self.act(x) |
|
x = self.downsample(x) |
|
return x |
|
|
|
|
|
class AMPBlock(nn.Sequential): |
|
def __init__(self, channels, *, kernel_size=3, dilations=(1, 3, 5)): |
|
super().__init__( |
|
*(self._make_layer(channels, kernel_size, d) for d in dilations) |
|
) |
|
|
|
def _make_layer(self, channels, kernel_size, dilation): |
|
return nn.Sequential( |
|
weight_norm( |
|
nn.Conv1d( |
|
channels, channels, kernel_size, dilation=dilation, padding="same" |
|
) |
|
), |
|
UpActDown(act=SnakeBeta(channels)), |
|
weight_norm(nn.Conv1d(channels, channels, kernel_size, padding="same")), |
|
) |
|
|
|
def forward(self, x): |
|
return x + super().forward(x) |
|
|
