Ref2VA/audio_vae/dac_activations.py
| 1 | # SPDX-License-Identifier: MIT |
| 2 | # Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. |
| 3 | |
| 4 | import torch |
| 5 | from torch import nn |
| 6 | from torch.nn import Parameter |
| 7 | |
| 8 | |
| 9 | @torch.jit.script |
| 10 | def snakebeta(x, alpha, beta): |
| 11 | shape = x.shape |
| 12 | x = x.reshape(shape[0], shape[1], -1) |
| 13 | x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) |
| 14 | x = x.reshape(shape) |
| 15 | return x |
| 16 | |
| 17 | |
| 18 | class SnakeBeta(nn.Module): |
| 19 | def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): |
| 20 | """ |
| 21 | Initialization. |
| 22 | INPUT: |
| 23 | - in_features: shape of the input |
| 24 | - alpha - trainable parameter that controls frequency |
| 25 | - beta - trainable parameter that controls magnitude |
| 26 | alpha is initialized to 1 by default, higher values = higher-frequency. |
| 27 | beta is initialized to 1 by default, higher values = higher-magnitude. |
| 28 | alpha will be trained along with the rest of your model. |
| 29 | """ |
| 30 | super(SnakeBeta, self).__init__() |
| 31 | self.in_features = in_features |
| 32 | |
| 33 | # Initialize alpha |
| 34 | self.alpha_logscale = alpha_logscale |
| 35 | if self.alpha_logscale: # Log scale alphas initialized to zeros |
| 36 | self.alpha = Parameter(torch.zeros(in_features) * alpha) |
| 37 | self.beta = Parameter(torch.zeros(in_features) * alpha) |
| 38 | else: # Linear scale alphas initialized to ones |
| 39 | self.alpha = Parameter(torch.ones(in_features) * alpha) |
| 40 | self.beta = Parameter(torch.ones(in_features) * alpha) |
| 41 | |
| 42 | self.alpha.requires_grad = alpha_trainable |
| 43 | self.beta.requires_grad = alpha_trainable |
| 44 | |
| 45 | self.no_div_by_zero = 0.000000001 |
| 46 | |
| 47 | def forward(self, x): |
| 48 | """ |
| 49 | Forward pass of the function. |
| 50 | Applies the function to the input elementwise. |
| 51 | SnakeBeta := x + 1/b * sin^2 (xa) |
| 52 | """ |
| 53 | alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] |
| 54 | beta = self.beta.unsqueeze(0).unsqueeze(-1) |
| 55 | if self.alpha_logscale: |
| 56 | alpha = torch.exp(alpha) |
| 57 | beta = torch.exp(beta) |
| 58 | x = snakebeta(x, alpha, beta) |
| 59 | |
| 60 | return x |
| 61 | |