File size: 2,255 Bytes
259eeac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import torch
import numpy as np
from torch import nn
from .module import WNConv1d, EncoderBlock
from .alias_free_torch import Activation1d
from . import activations
def init_weights(m):
if isinstance(m, nn.Conv1d):
nn.init.trunc_normal_(m.weight, std=0.02)
nn.init.constant_(m.bias, 0)
class CodecEncoder(nn.Module):
def __init__(
self,
ngf=48,
up_ratios=[2, 2, 4, 4, 5],
dilations=(1, 3, 9),
hidden_dim=1024,
depth=12,
heads=12,
pos_meb_dim=64,
):
super().__init__()
self.hop_length = np.prod(up_ratios)
self.ngf = ngf
self.up_ratios = up_ratios
d_model = ngf
self.conv_blocks = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
for i, stride in enumerate(up_ratios):
d_model *= 2
self.conv_blocks += [
EncoderBlock(d_model, stride=stride, dilations=dilations)
]
self.conv_blocks = nn.Sequential(*self.conv_blocks)
self.conv_final_block = [
Activation1d(
activation=activations.SnakeBeta(d_model, alpha_logscale=True)
),
WNConv1d(d_model, hidden_dim, kernel_size=3, padding=1),
]
self.conv_final_block = nn.Sequential(*self.conv_final_block)
self.reset_parameters()
def forward(self, x):
x = self.conv_blocks(x)
x = self.conv_final_block(x)
x = x.permute(0, 2, 1)
return x
def inference(self, x):
return self.block(x)
def remove_weight_norm(self):
"""Remove weight normalization module from all of the layers."""
def _remove_weight_norm(m):
try:
torch.nn.utils.remove_weight_norm(m)
except ValueError: # this module didn't have weight norm
return
self.apply(_remove_weight_norm)
def apply_weight_norm(self):
"""Apply weight normalization module from all of the layers."""
def _apply_weight_norm(m):
if isinstance(m, nn.Conv1d):
torch.nn.utils.weight_norm(m)
self.apply(_apply_weight_norm)
def reset_parameters(self):
self.apply(init_weights)
|