mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-10 10:32:17 +00:00
369 lines
13 KiB
Python
369 lines
13 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file at https://github.com/facebookresearch/encodec/blob/main/LICENSE
|
|
"""Encodec SEANet-based encoder and decoder implementation."""
|
|
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
from modules import SLSTM, SConv1d, SConvTranspose1d
|
|
|
|
|
|
class SEANetResnetBlock(nn.Module):
|
|
"""Residual block from SEANet model.
|
|
Args:
|
|
dim (int): Dimension of the input/output
|
|
kernel_sizes (list): List of kernel sizes for the convolutions.
|
|
dilations (list): List of dilations for the convolutions.
|
|
activation (str): Activation function.
|
|
activation_params (dict): Parameters to provide to the activation function
|
|
norm (str): Normalization method.
|
|
norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
|
|
causal (bool): Whether to use fully causal convolution.
|
|
pad_mode (str): Padding mode for the convolutions.
|
|
compress (int): Reduced dimensionality in residual branches (from Demucs v3)
|
|
true_skip (bool): Whether to use true skip connection or a simple convolution as the skip connection.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
kernel_sizes: List[int] = [3, 1],
|
|
dilations: List[int] = [1, 1],
|
|
activation: str = "ELU",
|
|
activation_params: Dict = {"alpha": 1.0},
|
|
norm: str = "weight_norm",
|
|
norm_params: Dict[str, Any] = {},
|
|
causal: bool = False,
|
|
pad_mode: str = "reflect",
|
|
compress: int = 2,
|
|
true_skip: bool = True,
|
|
):
|
|
super().__init__()
|
|
assert len(kernel_sizes) == len(
|
|
dilations
|
|
), "Number of kernel sizes should match number of dilations"
|
|
act = getattr(nn, activation)
|
|
hidden = dim // compress
|
|
block = []
|
|
for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
|
|
in_chs = dim if i == 0 else hidden
|
|
out_chs = dim if i == len(kernel_sizes) - 1 else hidden
|
|
block += [
|
|
act(**activation_params),
|
|
SConv1d(
|
|
in_chs,
|
|
out_chs,
|
|
kernel_size=kernel_size,
|
|
dilation=dilation,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
),
|
|
]
|
|
self.block = nn.Sequential(*block)
|
|
self.shortcut: nn.Module
|
|
if true_skip:
|
|
self.shortcut = nn.Identity()
|
|
else:
|
|
self.shortcut = SConv1d(
|
|
dim,
|
|
dim,
|
|
kernel_size=1,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.shortcut(x) + self.block(x)
|
|
|
|
|
|
class SEANetEncoder(nn.Module):
|
|
"""SEANet encoder.
|
|
Args:
|
|
channels (int): Audio channels.
|
|
dimension (int): Intermediate representation dimension.
|
|
n_filters (int): Base width for the model.
|
|
n_residual_layers (int): nb of residual layers.
|
|
ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
|
|
upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
|
|
that must match the decoder order
|
|
activation (str): Activation function.
|
|
activation_params (dict): Parameters to provide to the activation function
|
|
norm (str): Normalization method.
|
|
norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
|
|
kernel_size (int): Kernel size for the initial convolution.
|
|
last_kernel_size (int): Kernel size for the initial convolution.
|
|
residual_kernel_size (int): Kernel size for the residual layers.
|
|
dilation_base (int): How much to increase the dilation with each layer.
|
|
causal (bool): Whether to use fully causal convolution.
|
|
pad_mode (str): Padding mode for the convolutions.
|
|
true_skip (bool): Whether to use true skip connection or a simple
|
|
(streamable) convolution as the skip connection in the residual network blocks.
|
|
compress (int): Reduced dimensionality in residual branches (from Demucs v3).
|
|
lstm (int): Number of LSTM layers at the end of the encoder.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int = 1,
|
|
dimension: int = 128,
|
|
n_filters: int = 32,
|
|
n_residual_layers: int = 1,
|
|
ratios: List[int] = [8, 5, 4, 2],
|
|
activation: str = "ELU",
|
|
activation_params: dict = {"alpha": 1.0},
|
|
norm: str = "weight_norm",
|
|
norm_params: Dict[str, Any] = {},
|
|
kernel_size: int = 7,
|
|
last_kernel_size: int = 7,
|
|
residual_kernel_size: int = 3,
|
|
dilation_base: int = 2,
|
|
causal: bool = False,
|
|
pad_mode: str = "reflect",
|
|
true_skip: bool = False,
|
|
compress: int = 2,
|
|
lstm: int = 2,
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.dimension = dimension
|
|
self.n_filters = n_filters
|
|
self.ratios = list(reversed(ratios))
|
|
del ratios
|
|
self.n_residual_layers = n_residual_layers
|
|
self.hop_length = np.prod(self.ratios) # 计算乘积
|
|
|
|
act = getattr(nn, activation)
|
|
mult = 1
|
|
model: List[nn.Module] = [
|
|
SConv1d(
|
|
channels,
|
|
mult * n_filters,
|
|
kernel_size,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
)
|
|
]
|
|
# Downsample to raw audio scale
|
|
for i, ratio in enumerate(self.ratios):
|
|
# Add residual layers
|
|
for j in range(n_residual_layers):
|
|
model += [
|
|
SEANetResnetBlock(
|
|
mult * n_filters,
|
|
kernel_sizes=[residual_kernel_size, 1],
|
|
dilations=[dilation_base**j, 1],
|
|
norm=norm,
|
|
norm_params=norm_params,
|
|
activation=activation,
|
|
activation_params=activation_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
compress=compress,
|
|
true_skip=true_skip,
|
|
)
|
|
]
|
|
|
|
# Add downsampling layers
|
|
model += [
|
|
act(**activation_params),
|
|
SConv1d(
|
|
mult * n_filters,
|
|
mult * n_filters * 2,
|
|
kernel_size=ratio * 2,
|
|
stride=ratio,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
),
|
|
]
|
|
mult *= 2
|
|
|
|
if lstm:
|
|
model += [SLSTM(mult * n_filters, num_layers=lstm)]
|
|
|
|
model += [
|
|
act(**activation_params),
|
|
SConv1d(
|
|
mult * n_filters,
|
|
dimension,
|
|
last_kernel_size,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
),
|
|
]
|
|
|
|
self.model = nn.Sequential(*model)
|
|
|
|
def forward(self, x):
|
|
return self.model(x)
|
|
|
|
|
|
class SEANetDecoder(nn.Module):
|
|
"""SEANet decoder.
|
|
Args:
|
|
channels (int): Audio channels.
|
|
dimension (int): Intermediate representation dimension.
|
|
n_filters (int): Base width for the model.
|
|
n_residual_layers (int): nb of residual layers.
|
|
ratios (Sequence[int]): kernel size and stride ratios
|
|
activation (str): Activation function.
|
|
activation_params (dict): Parameters to provide to the activation function
|
|
final_activation (str): Final activation function after all convolutions.
|
|
final_activation_params (dict): Parameters to provide to the activation function
|
|
norm (str): Normalization method.
|
|
norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
|
|
kernel_size (int): Kernel size for the initial convolution.
|
|
last_kernel_size (int): Kernel size for the initial convolution.
|
|
residual_kernel_size (int): Kernel size for the residual layers.
|
|
dilation_base (int): How much to increase the dilation with each layer.
|
|
causal (bool): Whether to use fully causal convolution.
|
|
pad_mode (str): Padding mode for the convolutions.
|
|
true_skip (bool): Whether to use true skip connection or a simple
|
|
(streamable) convolution as the skip connection in the residual network blocks.
|
|
compress (int): Reduced dimensionality in residual branches (from Demucs v3).
|
|
lstm (int): Number of LSTM layers at the end of the encoder.
|
|
trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
|
|
If equal to 1.0, it means that all the trimming is done at the right.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int = 1,
|
|
dimension: int = 128,
|
|
n_filters: int = 32,
|
|
n_residual_layers: int = 1,
|
|
ratios: List[int] = [8, 5, 4, 2],
|
|
activation: str = "ELU",
|
|
activation_params: dict = {"alpha": 1.0},
|
|
final_activation: Optional[str] = None,
|
|
final_activation_params: Optional[dict] = None,
|
|
norm: str = "weight_norm",
|
|
norm_params: Dict[str, Any] = {},
|
|
kernel_size: int = 7,
|
|
last_kernel_size: int = 7,
|
|
residual_kernel_size: int = 3,
|
|
dilation_base: int = 2,
|
|
causal: bool = False,
|
|
pad_mode: str = "reflect",
|
|
true_skip: bool = False,
|
|
compress: int = 2,
|
|
lstm: int = 2,
|
|
trim_right_ratio: float = 1.0,
|
|
):
|
|
super().__init__()
|
|
self.dimension = dimension
|
|
self.channels = channels
|
|
self.n_filters = n_filters
|
|
self.ratios = ratios
|
|
del ratios
|
|
self.n_residual_layers = n_residual_layers
|
|
self.hop_length = np.prod(self.ratios)
|
|
|
|
act = getattr(nn, activation)
|
|
mult = int(2 ** len(self.ratios))
|
|
model: List[nn.Module] = [
|
|
SConv1d(
|
|
dimension,
|
|
mult * n_filters,
|
|
kernel_size,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
)
|
|
]
|
|
|
|
if lstm:
|
|
model += [SLSTM(mult * n_filters, num_layers=lstm)]
|
|
|
|
# Upsample to raw audio scale
|
|
for i, ratio in enumerate(self.ratios):
|
|
# Add upsampling layers
|
|
model += [
|
|
act(**activation_params),
|
|
SConvTranspose1d(
|
|
mult * n_filters,
|
|
mult * n_filters // 2,
|
|
kernel_size=ratio * 2,
|
|
stride=ratio,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
trim_right_ratio=trim_right_ratio,
|
|
),
|
|
]
|
|
# Add residual layers
|
|
for j in range(n_residual_layers):
|
|
model += [
|
|
SEANetResnetBlock(
|
|
mult * n_filters // 2,
|
|
kernel_sizes=[residual_kernel_size, 1],
|
|
dilations=[dilation_base**j, 1],
|
|
activation=activation,
|
|
activation_params=activation_params,
|
|
norm=norm,
|
|
norm_params=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
compress=compress,
|
|
true_skip=true_skip,
|
|
)
|
|
]
|
|
|
|
mult //= 2
|
|
|
|
# Add final layers
|
|
model += [
|
|
act(**activation_params),
|
|
SConv1d(
|
|
n_filters,
|
|
channels,
|
|
last_kernel_size,
|
|
norm=norm,
|
|
norm_kwargs=norm_params,
|
|
causal=causal,
|
|
pad_mode=pad_mode,
|
|
),
|
|
]
|
|
# Add optional final activation to decoder (eg. tanh)
|
|
if final_activation is not None:
|
|
final_act = getattr(nn, final_activation)
|
|
final_activation_params = final_activation_params or {}
|
|
model += [final_act(**final_activation_params)]
|
|
self.model = nn.Sequential(*model)
|
|
|
|
def forward(self, z):
|
|
y = self.model(z)
|
|
return y
|
|
|
|
|
|
def test():
|
|
import torch
|
|
|
|
encoder = SEANetEncoder()
|
|
decoder = SEANetDecoder()
|
|
x = torch.randn(1, 1, 24000)
|
|
z = encoder(x)
|
|
print("z ", z.shape)
|
|
assert 1 == 2
|
|
assert list(z.shape) == [1, 128, 75], z.shape
|
|
y = decoder(z)
|
|
assert y.shape == x.shape, (x.shape, y.shape)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test()
|