Clicko777/RVC_HFv2

1# File under the MIT license, see https://github.com/adefossez/julius/LICENSE for details.2# Author: adefossez, 20203 4"""5Implementation of a FFT based 1D convolution in PyTorch.6While FFT is used in CUDNN for small kernel sizes, it is not the case for long ones, e.g. 512.7This module implements efficient FFT based convolutions for such convolutions. A typical8application is for evaluationg FIR filters with a long receptive field, typically9evaluated with a stride of 1.10"""11from typing import Optional12 13import torch14try:15    import torch.fft as new_fft16except ImportError:17    new_fft = None  # type: ignore18from torch.nn import functional as F19 20from .core import pad_to, unfold21from .utils import simple_repr22 23 24# This is quite verbose, but sadly needed to make TorchScript happy.25def _new_rfft(x: torch.Tensor):26    z = new_fft.rfft(x, dim=-1)27    return torch.view_as_real(z)28 29 30def _old_rfft(x: torch.Tensor):31    return torch.rfft(x, 1)  # type: ignore32 33 34def _old_irfft(x: torch.Tensor, length: int):35    result = torch.irfft(x, 1, signal_sizes=(length,))  # type: ignore36    return result37 38 39def _new_irfft(x: torch.Tensor, length: int):40    x = torch.view_as_complex(x)41    return new_fft.irfft(x, length, dim=-1)42 43 44if new_fft is None:45    _rfft = _old_rfft46    _irfft = _old_irfft47else:48    _rfft = _new_rfft49    _irfft = _new_irfft50 51 52def _compl_mul_conjugate(a: torch.Tensor, b: torch.Tensor):53    """54    Given a and b two tensors of dimension 455    with the last dimension being the real and imaginary part,56    returns a multiplied by the conjugate of b, the multiplication57    being with respect to the second dimension.58 59    """60    # PyTorch 1.7 supports complex number, but not for all operations.61    # Once the support is widespread, this can likely go away.62 63    op = "bcft,dct->bdft"64    return torch.stack([65        torch.einsum(op, a[..., 0], b[..., 0]) + torch.einsum(op, a[..., 1], b[..., 1]),66        torch.einsum(op, a[..., 1], b[..., 0]) - torch.einsum(op, a[..., 0], b[..., 1])67    ],68                       dim=-1)69 70 71def fft_conv1d(72        input: torch.Tensor, weight: torch.Tensor,73        bias: Optional[torch.Tensor] = None, stride: int = 1, padding: int = 0,74        block_ratio: float = 5):75    """76    Same as `torch.nn.functional.conv1d` but using FFT for the convolution.77    Please check PyTorch documentation for more information.78 79    Args:80        input (Tensor): input signal of shape `[B, C, T]`.81        weight (Tensor): weight of the convolution `[D, C, K]` with `D` the number82            of output channels.83        bias (Tensor or None): if not None, bias term for the convolution.84        stride (int): stride of convolution.85        padding (int): padding to apply to the input.86        block_ratio (float): can be tuned for speed. The input is splitted in chunks87            with a size of `int(block_ratio * kernel_size)`.88 89    Shape:90 91        - Inputs: `input` is `[B, C, T]`, `weight` is `[D, C, K]` and bias is `[D]`.92        - Output: `(*, T)`93 94 95    ..note::96        This function is faster than `torch.nn.functional.conv1d` only in specific cases.97        Typically, the kernel size should be of the order of 256 to see any real gain,98        for a stride of 1.99 100    ..Warning::101        Dilation and groups are not supported at the moment. This function might use102        more memory than the default Conv1d implementation.103    """104    input = F.pad(input, (padding, padding))105    batch, channels, length = input.shape106    out_channels, _, kernel_size = weight.shape107 108    if length < kernel_size:109        raise RuntimeError(f"Input should be at least as large as the kernel size {kernel_size}, "110                           f"but it is only {length} samples long.")111    if block_ratio < 1:112        raise RuntimeError("Block ratio must be greater than 1.")113 114    # We are going to process the input blocks by blocks, as for some reason it is faster115    # and less memory intensive (I think the culprit is `torch.einsum`.116    block_size: int = min(int(kernel_size * block_ratio), length)117    fold_stride = block_size - kernel_size + 1118    weight = pad_to(weight, block_size)119    weight_z = _rfft(weight)120 121    # We pad the input and get the different frames, on which122    frames = unfold(input, block_size, fold_stride)123 124    frames_z = _rfft(frames)125    out_z = _compl_mul_conjugate(frames_z, weight_z)126    out = _irfft(out_z, block_size)127    # The last bit is invalid, because FFT will do a circular convolution.128    out = out[..., :-kernel_size + 1]129    out = out.reshape(batch, out_channels, -1)130    out = out[..., ::stride]131    target_length = (length - kernel_size) // stride + 1132    out = out[..., :target_length]133    if bias is not None:134        out += bias[:, None]135    return out136 137 138class FFTConv1d(torch.nn.Module):139    """140    Same as `torch.nn.Conv1d` but based on `fft_conv1d`.141    Please check PyTorch documentation for more information.142 143    Args:144        in_channels (int): number of input channels.145        out_channels (int): number of output channels.146        kernel_size (int): kernel size of convolution.147        stride (int): stride of convolution.148        padding (int): padding to apply to the input.149        bias (bool): if True, use a bias term.150 151    ..note::152        This module is faster than `torch.nn.Conv1d` only in specific cases.153        Typically, `kernel_size` should be of the order of 256 to see any real gain,154        for a stride of 1.155 156    ..warning::157        Dilation and groups are not supported at the moment. This module might use158        more memory than the default Conv1d implementation.159 160    >>> fftconv = FFTConv1d(12, 24, 128, 4)161    >>> x = torch.randn(4, 12, 1024)162    >>> print(list(fftconv(x).shape))163    [4, 24, 225]164    """165    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,166                 stride: int = 1, padding: int = 0, bias: bool = True):167        super().__init__()168        self.in_channels = in_channels169        self.out_channels = out_channels170        self.kernel_size = kernel_size171        self.stride = stride172        self.padding = padding173 174        conv = torch.nn.Conv1d(in_channels, out_channels, kernel_size, bias=bias)175        self.weight = conv.weight176        self.bias = conv.bias177 178    def forward(self, input: torch.Tensor):179        return fft_conv1d(180            input, self.weight, self.bias, self.stride, self.padding)181 182    def __repr__(self):183        return simple_repr(self, overrides={"bias": self.bias is not None})184