scrapegoat/Neural-Audio-Codec

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""Convolutional layers wrappers and utilities."""8 9import math10import typing as tp11import warnings12 13import torch14from torch import nn15from torch.nn import functional as F16from torch.nn.utils import spectral_norm, weight_norm17 18from .norm import ConvLayerNorm19 20 21CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',22                                 'time_layer_norm', 'layer_norm', 'time_group_norm'])23 24 25def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:26    assert norm in CONV_NORMALIZATIONS27    if norm == 'weight_norm':28        return weight_norm(module)29    elif norm == 'spectral_norm':30        return spectral_norm(module)31    else:32        # We already check was in CONV_NORMALIZATION, so any other choice33        # doesn't need reparametrization.34        return module35 36 37def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:38    """Return the proper normalization module. If causal is True, this will ensure the returned39    module is causal, or return an error if the normalization doesn't support causal evaluation.40    """41    assert norm in CONV_NORMALIZATIONS42    if norm == 'layer_norm':43        assert isinstance(module, nn.modules.conv._ConvNd)44        return ConvLayerNorm(module.out_channels, **norm_kwargs)45    elif norm == 'time_group_norm':46        if causal:47            raise ValueError("GroupNorm doesn't support causal evaluation.")48        assert isinstance(module, nn.modules.conv._ConvNd)49        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)50    else:51        return nn.Identity()52 53 54def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,55                                 padding_total: int = 0) -> int:56    """See `pad_for_conv1d`.57    """58    length = x.shape[-1]59    n_frames = (length - kernel_size + padding_total) / stride + 160    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)61    return ideal_length - length62 63 64def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):65    """Pad for a convolution to make sure that the last window is full.66    Extra padding is added at the end. This is required to ensure that we can rebuild67    an output of the same length, as otherwise, even with padding, some time steps68    might get removed.69    For instance, with total padding = 4, kernel size = 4, stride = 2:70        0 0 1 2 3 4 5 0 0   # (0s are padding)71        1   2   3           # (output frames of a convolution, last 0 is never used)72        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)73            1 2 3 4         # once you removed padding, we are missing one time step !74    """75    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)76    return F.pad(x, (0, extra_padding))77 78 79def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):80    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.81    If this is the case, we insert extra 0 padding to the right before the reflection happen.82    """83    length = x.shape[-1]84    padding_left, padding_right = paddings85    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)86    if mode == 'reflect':87        max_pad = max(padding_left, padding_right)88        extra_pad = 089        if length <= max_pad:90            extra_pad = max_pad - length + 191            x = F.pad(x, (0, extra_pad))92        padded = F.pad(x, paddings, mode, value)93        end = padded.shape[-1] - extra_pad94        return padded[..., :end]95    else:96        return F.pad(x, paddings, mode, value)97 98 99def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):100    """Remove padding from x, handling properly zero padding. Only for 1d!"""101    padding_left, padding_right = paddings102    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)103    assert (padding_left + padding_right) <= x.shape[-1]104    end = x.shape[-1] - padding_right105    return x[..., padding_left: end]106 107 108class NormConv1d(nn.Module):109    """Wrapper around Conv1d and normalization applied to this conv110    to provide a uniform interface across normalization approaches.111    """112    def __init__(self, *args, causal: bool = False, norm: str = 'none',113                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):114        super().__init__()115        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)116        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)117        self.norm_type = norm118 119    def forward(self, x):120        x = self.conv(x)121        x = self.norm(x)122        return x123 124 125class NormConv2d(nn.Module):126    """Wrapper around Conv2d and normalization applied to this conv127    to provide a uniform interface across normalization approaches.128    """129    def __init__(self, *args, norm: str = 'none',130                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):131        super().__init__()132        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)133        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)134        self.norm_type = norm135 136    def forward(self, x):137        x = self.conv(x)138        x = self.norm(x)139        return x140 141 142class NormConvTranspose1d(nn.Module):143    """Wrapper around ConvTranspose1d and normalization applied to this conv144    to provide a uniform interface across normalization approaches.145    """146    def __init__(self, *args, causal: bool = False, norm: str = 'none',147                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):148        super().__init__()149        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)150        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)151        self.norm_type = norm152 153    def forward(self, x):154        x = self.convtr(x)155        x = self.norm(x)156        return x157 158 159class NormConvTranspose2d(nn.Module):160    """Wrapper around ConvTranspose2d and normalization applied to this conv161    to provide a uniform interface across normalization approaches.162    """163    def __init__(self, *args, norm: str = 'none',164                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):165        super().__init__()166        self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)167        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)168 169    def forward(self, x):170        x = self.convtr(x)171        x = self.norm(x)172        return x173 174 175class SConv1d(nn.Module):176    """Conv1d with some builtin handling of asymmetric or causal padding177    and normalization.178    """179    def __init__(self, in_channels: int, out_channels: int,180                 kernel_size: int, stride: int = 1, dilation: int = 1,181                 groups: int = 1, bias: bool = True, causal: bool = False,182                 norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},183                 pad_mode: str = 'reflect'):184        super().__init__()185        # warn user on unusual setup between dilation and stride186        if stride > 1 and dilation > 1:187            warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1'188                          f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).')189        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,190                               dilation=dilation, groups=groups, bias=bias, causal=causal,191                               norm=norm, norm_kwargs=norm_kwargs)192        self.causal = causal193        self.pad_mode = pad_mode194 195    def forward(self, x):196        B, C, T = x.shape197        kernel_size = self.conv.conv.kernel_size[0]198        stride = self.conv.conv.stride[0]199        dilation = self.conv.conv.dilation[0]200        padding_total = (kernel_size - 1) * dilation - (stride - 1)201        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)202        if self.causal:203            # Left padding for causal204            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)205        else:206            # Asymmetric padding required for odd strides207            padding_right = padding_total // 2208            padding_left = padding_total - padding_right209            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)210        return self.conv(x)211 212 213class SConvTranspose1d(nn.Module):214    """ConvTranspose1d with some builtin handling of asymmetric or causal padding215    and normalization.216    """217    def __init__(self, in_channels: int, out_channels: int,218                 kernel_size: int, stride: int = 1, causal: bool = False,219                 norm: str = 'none', trim_right_ratio: float = 1.,220                 norm_kwargs: tp.Dict[str, tp.Any] = {}):221        super().__init__()222        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,223                                          causal=causal, norm=norm, norm_kwargs=norm_kwargs)224        self.causal = causal225        self.trim_right_ratio = trim_right_ratio226        assert self.causal or self.trim_right_ratio == 1., \227            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"228        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.229 230    def forward(self, x):231        kernel_size = self.convtr.convtr.kernel_size[0]232        stride = self.convtr.convtr.stride[0]233        padding_total = kernel_size - stride234 235        y = self.convtr(x)236 237        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be238        # removed at the very end, when keeping only the right length for the output,239        # as removing it here would require also passing the length at the matching layer240        # in the encoder.241        if self.causal:242            # Trim the padding on the right according to the specified ratio243            # if trim_right_ratio = 1.0, trim everything from right244            padding_right = math.ceil(padding_total * self.trim_right_ratio)245            padding_left = padding_total - padding_right246            y = unpad1d(y, (padding_left, padding_right))247        else:248            # Asymmetric padding required for odd strides249            padding_right = padding_total // 2250            padding_left = padding_total - padding_right251            y = unpad1d(y, (padding_left, padding_right))252        return y253