mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-10 10:32:17 +00:00
335 lines
11 KiB
Python
335 lines
11 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
|
|
"""Convolutional layers wrappers and utilities."""
|
|
import logging
|
|
import math
|
|
from typing import Any, Dict, Tuple
|
|
|
|
from torch import Tensor, nn
|
|
from torch.nn import functional as F
|
|
from torch.nn.utils import spectral_norm, weight_norm
|
|
|
|
from .norm import ConvLayerNorm
|
|
|
|
CONV_NORMALIZATIONS = frozenset(
|
|
[
|
|
"none",
|
|
"weight_norm",
|
|
"spectral_norm",
|
|
"time_layer_norm",
|
|
"layer_norm",
|
|
"time_group_norm",
|
|
]
|
|
)
|
|
|
|
|
|
def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module:
|
|
assert norm in CONV_NORMALIZATIONS
|
|
if norm == "weight_norm":
|
|
return weight_norm(module)
|
|
elif norm == "spectral_norm":
|
|
return spectral_norm(module)
|
|
else:
|
|
# We already check was in CONV_NORMALIZATION, so any other choice
|
|
# doesn't need reparametrization.
|
|
return module
|
|
|
|
|
|
def get_norm_module(
|
|
module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs
|
|
) -> nn.Module:
|
|
"""Return the proper normalization module. If causal is True, this will ensure the returned
|
|
module is causal, or return an error if the normalization doesn't support causal evaluation.
|
|
"""
|
|
assert norm in CONV_NORMALIZATIONS
|
|
if norm == "layer_norm":
|
|
assert isinstance(module, nn.modules.conv._ConvNd)
|
|
return ConvLayerNorm(module.out_channels, **norm_kwargs)
|
|
elif norm == "time_group_norm":
|
|
if causal:
|
|
raise ValueError("GroupNorm doesn't support causal evaluation.")
|
|
assert isinstance(module, nn.modules.conv._ConvNd)
|
|
return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
|
|
else:
|
|
return nn.Identity()
|
|
|
|
|
|
def get_extra_padding_for_conv1d(
|
|
x: Tensor, kernel_size: int, stride: int, padding_total: int = 0
|
|
) -> int:
|
|
"""See `pad_for_conv1d`."""
|
|
length = x.shape[-1]
|
|
n_frames = (length - kernel_size + padding_total) / stride + 1
|
|
ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
|
|
return ideal_length - length
|
|
|
|
|
|
def pad_for_conv1d(x: Tensor, kernel_size: int, stride: int, padding_total: int = 0):
|
|
"""Pad for a convolution to make sure that the last window is full.
|
|
Extra padding is added at the end. This is required to ensure that we can rebuild
|
|
an output of the same length, as otherwise, even with padding, some time steps
|
|
might get removed.
|
|
For instance, with total padding = 4, kernel size = 4, stride = 2:
|
|
0 0 1 2 3 4 5 0 0 # (0s are padding)
|
|
1 2 3 # (output frames of a convolution, last 0 is never used)
|
|
0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
|
|
1 2 3 4 # once you removed padding, we are missing one time step !
|
|
"""
|
|
extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
|
|
return F.pad(x, (0, extra_padding))
|
|
|
|
|
|
def pad1d(
|
|
x: Tensor,
|
|
paddings: Tuple[int, int],
|
|
mode: str = "zero",
|
|
value: float = 0.0,
|
|
):
|
|
"""Tiny wrapper around F.pad, just to allow for reflect padding on small input.
|
|
If this is the case, we insert extra 0 padding to the right before the reflection happen.
|
|
"""
|
|
length = x.shape[-1]
|
|
padding_left, padding_right = paddings
|
|
assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
|
|
if mode == "reflect":
|
|
max_pad = max(padding_left, padding_right)
|
|
extra_pad = 0
|
|
if length <= max_pad:
|
|
extra_pad = max_pad - length + 1
|
|
x = F.pad(x, (0, extra_pad))
|
|
padded = F.pad(x, paddings, mode, value)
|
|
end = padded.shape[-1] - extra_pad
|
|
return padded[..., :end]
|
|
else:
|
|
return F.pad(x, paddings, mode, value)
|
|
|
|
|
|
def unpad1d(x: Tensor, paddings: Tuple[int, int]):
|
|
"""Remove padding from x, handling properly zero padding. Only for 1d!"""
|
|
padding_left, padding_right = paddings
|
|
assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
|
|
assert (padding_left + padding_right) <= x.shape[-1]
|
|
end = x.shape[-1] - padding_right
|
|
return x[..., padding_left:end]
|
|
|
|
|
|
class NormConv1d(nn.Module):
|
|
"""Wrapper around Conv1d and normalization applied to this conv
|
|
to provide a uniform interface across normalization approaches.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
causal: bool = False,
|
|
norm: str = "none",
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
|
|
self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
|
|
self.norm_type = norm
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.norm(x)
|
|
return x
|
|
|
|
|
|
class NormConv2d(nn.Module):
|
|
"""Wrapper around Conv2d and normalization applied to this conv
|
|
to provide a uniform interface across normalization approaches.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
norm: str = "none",
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
|
|
self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
|
|
self.norm_type = norm
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.norm(x)
|
|
return x
|
|
|
|
|
|
class NormConvTranspose1d(nn.Module):
|
|
"""Wrapper around ConvTranspose1d and normalization applied to this conv
|
|
to provide a uniform interface across normalization approaches.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
causal: bool = False,
|
|
norm: str = "none",
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.convtr = apply_parametrization_norm(
|
|
nn.ConvTranspose1d(*args, **kwargs), norm
|
|
)
|
|
self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
|
|
self.norm_type = norm
|
|
|
|
def forward(self, x):
|
|
x = self.convtr(x)
|
|
x = self.norm(x)
|
|
return x
|
|
|
|
|
|
class NormConvTranspose2d(nn.Module):
|
|
"""Wrapper around ConvTranspose2d and normalization applied to this conv
|
|
to provide a uniform interface across normalization approaches.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
norm: str = "none",
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.convtr = apply_parametrization_norm(
|
|
nn.ConvTranspose2d(*args, **kwargs), norm
|
|
)
|
|
self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
|
|
|
|
def forward(self, x):
|
|
x = self.convtr(x)
|
|
x = self.norm(x)
|
|
return x
|
|
|
|
|
|
class SConv1d(nn.Module):
|
|
"""Conv1d with some builtin handling of asymmetric or causal padding
|
|
and normalization.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
stride: int = 1,
|
|
dilation: int = 1,
|
|
groups: int = 1,
|
|
bias: bool = True,
|
|
causal: bool = False,
|
|
norm: str = "none",
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
pad_mode: str = "reflect",
|
|
):
|
|
super().__init__()
|
|
# warn user on unusual setup between dilation and stride
|
|
if stride > 1 and dilation > 1:
|
|
logging.warning(
|
|
"SConv1d has been initialized with stride > 1 and dilation > 1"
|
|
f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
|
|
)
|
|
self.conv = NormConv1d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride,
|
|
dilation=dilation,
|
|
groups=groups,
|
|
bias=bias,
|
|
causal=causal,
|
|
norm=norm,
|
|
norm_kwargs=norm_kwargs,
|
|
)
|
|
self.causal = causal
|
|
self.pad_mode = pad_mode
|
|
|
|
def forward(self, x):
|
|
B, C, T = x.shape
|
|
kernel_size = self.conv.conv.kernel_size[0]
|
|
stride = self.conv.conv.stride[0]
|
|
dilation = self.conv.conv.dilation[0]
|
|
padding_total = (kernel_size - 1) * dilation - (stride - 1)
|
|
extra_padding = get_extra_padding_for_conv1d(
|
|
x, kernel_size, stride, padding_total
|
|
)
|
|
if self.causal:
|
|
# Left padding for causal
|
|
x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
|
|
else:
|
|
# Asymmetric padding required for odd strides
|
|
padding_right = padding_total // 2
|
|
padding_left = padding_total - padding_right
|
|
x = pad1d(
|
|
x, (padding_left, padding_right + extra_padding), mode=self.pad_mode
|
|
)
|
|
return self.conv(x)
|
|
|
|
|
|
class SConvTranspose1d(nn.Module):
|
|
"""ConvTranspose1d with some builtin handling of asymmetric or causal padding
|
|
and normalization.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
stride: int = 1,
|
|
causal: bool = False,
|
|
norm: str = "none",
|
|
trim_right_ratio: float = 1.0,
|
|
norm_kwargs: Dict[str, Any] = {},
|
|
):
|
|
super().__init__()
|
|
self.convtr = NormConvTranspose1d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride,
|
|
causal=causal,
|
|
norm=norm,
|
|
norm_kwargs=norm_kwargs,
|
|
)
|
|
self.causal = causal
|
|
self.trim_right_ratio = trim_right_ratio
|
|
assert (
|
|
self.causal or self.trim_right_ratio == 1.0
|
|
), "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
|
|
assert self.trim_right_ratio >= 0.0 and self.trim_right_ratio <= 1.0
|
|
|
|
def forward(self, x):
|
|
kernel_size = self.convtr.convtr.kernel_size[0]
|
|
stride = self.convtr.convtr.stride[0]
|
|
padding_total = kernel_size - stride
|
|
|
|
y = self.convtr(x)
|
|
|
|
# We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
|
|
# removed at the very end, when keeping only the right length for the output,
|
|
# as removing it here would require also passing the length at the matching layer
|
|
# in the encoder.
|
|
if self.causal:
|
|
# Trim the padding on the right according to the specified ratio
|
|
# if trim_right_ratio = 1.0, trim everything from right
|
|
padding_right = math.ceil(padding_total * self.trim_right_ratio)
|
|
padding_left = padding_total - padding_right
|
|
y = unpad1d(y, (padding_left, padding_right))
|
|
else:
|
|
# Asymmetric padding required for odd strides
|
|
padding_right = padding_total // 2
|
|
padding_left = padding_total - padding_right
|
|
y = unpad1d(y, (padding_left, padding_right))
|
|
return y
|