Remove previous DPTNet/DCUNet implementations

cfg/model/dptnet.yaml

@@ -10,12 +10,13 @@ model:
   network:
     _target_: remfx.models.DPTNetModel
     n_src: 1
-    bn_chan: 128
-    hid_size: 128
+    in_chan: 64
+    out_chan: 64
     chunk_size: 100
-    n_repeats: 6
+    n_repeats: 2
     fb_name: "free"
     kernel_size: 16
-    n_filters: 1
+    n_filters: 64
+    stride: 8
     sample_rate: ${sample_rate}
     num_bins: 1025

cfg/model/tcn.yaml

@@ -13,7 +13,7 @@ model:
     noutputs: 1
     nblocks: 20
     channel_growth: 0
-    channel_width: 32
+    channel_width: 64
     kernel_size: 7
     stack_size: 10
     dilation_growth: 2

remfx/datasets.py

@@ -295,7 +295,7 @@ class EffectDataset(Dataset):
 
         # Up to max_kept_effects
         if self.max_kept_effects != -1:
-            num_kept_effects = int(torch.rand(1).item() * (self.max_kept_effects)) + 1
+            num_kept_effects = int(torch.rand(1).item() * (self.max_kept_effects))
         else:
             num_kept_effects = len(self.effects_to_keep)
         effect_indices = effect_indices[:num_kept_effects]

remfx/dcunet.py

@@ -1,649 +0,0 @@
-# Adapted from https://github.com/AppleHolic/source_separation/tree/master/source_separation
-
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from torch.nn.init import calculate_gain
-from typing import Tuple
-from scipy.signal import get_window
-from librosa.util import pad_center
-from remfx.utils import single, concat_complex
-
-
-class ComplexConvBlock(nn.Module):
-    """
-    Convolution block
-    """
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        padding: int = 0,
-        layers: int = 4,
-        bn_func=nn.BatchNorm1d,
-        act_func=nn.LeakyReLU,
-        skip_res: bool = False,
-    ):
-        super().__init__()
-        # modules
-        self.blocks = nn.ModuleList()
-        self.skip_res = skip_res
-
-        for idx in range(layers):
-            in_ = in_channels if idx == 0 else out_channels
-            self.blocks.append(
-                nn.Sequential(
-                    *[
-                        bn_func(in_),
-                        act_func(),
-                        ComplexConv1d(in_, out_channels, kernel_size, padding=padding),
-                    ]
-                )
-            )
-
-    def forward(self, x: torch.tensor) -> torch.tensor:
-        temp = x
-        for idx, block in enumerate(self.blocks):
-            x = block(x)
-
-        if temp.size() != x.size() or self.skip_res:
-            return x
-        else:
-            return x + temp
-
-
-class SpectrogramUnet(nn.Module):
-    def __init__(
-        self,
-        spec_dim: int,
-        hidden_dim: int,
-        filter_len: int,
-        hop_len: int,
-        layers: int = 3,
-        block_layers: int = 3,
-        kernel_size: int = 5,
-        is_mask: bool = False,
-        norm: str = "bn",
-        act: str = "tanh",
-    ):
-        super().__init__()
-        self.layers = layers
-        self.is_mask = is_mask
-
-        # stft modules
-        self.stft = STFT(filter_len, hop_len)
-
-        if norm == "bn":
-            self.bn_func = nn.BatchNorm1d
-        elif norm == "ins":
-            self.bn_func = lambda x: nn.InstanceNorm1d(x, affine=True)
-        else:
-            raise NotImplementedError("{} is not implemented !".format(norm))
-
-        if act == "tanh":
-            self.act_func = nn.Tanh
-            self.act_out = nn.Tanh
-        elif act == "comp":
-            self.act_func = ComplexActLayer
-            self.act_out = lambda: ComplexActLayer(is_out=True)
-        else:
-            raise NotImplementedError("{} is not implemented !".format(act))
-
-        # prev conv
-        self.prev_conv = ComplexConv1d(spec_dim * 2, hidden_dim, 1)
-
-        # down
-        self.down = nn.ModuleList()
-        self.down_pool = nn.MaxPool1d(3, stride=2, padding=1)
-        for idx in range(self.layers):
-            block = ComplexConvBlock(
-                hidden_dim,
-                hidden_dim,
-                kernel_size=kernel_size,
-                padding=kernel_size // 2,
-                bn_func=self.bn_func,
-                act_func=self.act_func,
-                layers=block_layers,
-            )
-            self.down.append(block)
-
-        # up
-        self.up = nn.ModuleList()
-        for idx in range(self.layers):
-            in_c = hidden_dim if idx == 0 else hidden_dim * 2
-            self.up.append(
-                nn.Sequential(
-                    ComplexConvBlock(
-                        in_c,
-                        hidden_dim,
-                        kernel_size=kernel_size,
-                        padding=kernel_size // 2,
-                        bn_func=self.bn_func,
-                        act_func=self.act_func,
-                        layers=block_layers,
-                    ),
-                    self.bn_func(hidden_dim),
-                    self.act_func(),
-                    ComplexTransposedConv1d(
-                        hidden_dim, hidden_dim, kernel_size=2, stride=2
-                    ),
-                )
-            )
-
-        # out_conv
-        self.out_conv = nn.Sequential(
-            ComplexConvBlock(
-                hidden_dim * 2,
-                spec_dim * 2,
-                kernel_size=kernel_size,
-                padding=kernel_size // 2,
-                bn_func=self.bn_func,
-                act_func=self.act_func,
-            ),
-            self.bn_func(spec_dim * 2),
-            self.act_func(),
-        )
-
-        # refine conv
-        self.refine_conv = nn.Sequential(
-            ComplexConvBlock(
-                spec_dim * 4,
-                spec_dim * 2,
-                kernel_size=kernel_size,
-                padding=kernel_size // 2,
-                bn_func=self.bn_func,
-                act_func=self.act_func,
-            ),
-            self.bn_func(spec_dim * 2),
-            self.act_func(),
-        )
-
-    def log_stft(self, wav):
-        # stft
-        mag, phase = self.stft.transform(wav)
-        return torch.log(mag + 1), phase
-
-    def exp_istft(self, log_mag, phase):
-        # exp
-        mag = np.e**log_mag - 1
-        # istft
-        wav = self.stft.inverse(mag, phase)
-        return wav
-
-    def adjust_diff(self, x, target):
-        size_diff = target.size()[-1] - x.size()[-1]
-        assert size_diff >= 0
-        if size_diff > 0:
-            x = F.pad(
-                x.unsqueeze(1), (size_diff // 2, size_diff // 2), "reflect"
-            ).squeeze(1)
-        return x
-
-    def masking(self, mag, phase, origin_mag, origin_phase):
-        abs_mag = torch.abs(mag)
-        mag_mask = torch.tanh(abs_mag)
-        phase_mask = mag / abs_mag
-
-        # masking
-        mag = mag_mask * origin_mag
-        phase = phase_mask * (origin_phase + phase)
-        return mag, phase
-
-    def forward(self, wav):
-        # stft
-        origin_mag, origin_phase = self.log_stft(wav)
-        origin_x = torch.cat([origin_mag, origin_phase], dim=1)
-
-        # prev
-        x = self.prev_conv(origin_x)
-
-        # body
-        # down
-        down_cache = []
-        for idx, block in enumerate(self.down):
-            x = block(x)
-            down_cache.append(x)
-            x = self.down_pool(x)
-
-        # up
-        for idx, block in enumerate(self.up):
-            x = block(x)
-            res = F.interpolate(
-                down_cache[self.layers - (idx + 1)],
-                size=[x.size()[2]],
-                mode="linear",
-                align_corners=False,
-            )
-            x = concat_complex(x, res, dim=1)
-
-        # match spec dimension
-        x = self.out_conv(x)
-        if origin_mag.size(2) != x.size(2):
-            x = F.interpolate(
-                x, size=[origin_mag.size(2)], mode="linear", align_corners=False
-            )
-
-        # refine
-        x = self.refine_conv(concat_complex(x, origin_x))
-
-        def to_wav(stft):
-            mag, phase = stft.chunk(2, 1)
-            if self.is_mask:
-                mag, phase = self.masking(mag, phase, origin_mag, origin_phase)
-            out = self.exp_istft(mag, phase)
-            out = self.adjust_diff(out, wav)
-            return out
-
-        refine_wav = to_wav(x)
-
-        return refine_wav
-
-
-class RefineSpectrogramUnet(SpectrogramUnet):
-    def __init__(
-        self,
-        spec_dim: int,
-        hidden_dim: int,
-        filter_len: int,
-        hop_len: int,
-        layers: int = 4,
-        block_layers: int = 4,
-        kernel_size: int = 3,
-        is_mask: bool = True,
-        norm: str = "ins",
-        act: str = "comp",
-        refine_layers: int = 1,
-        add_spec_results: bool = False,
-    ):
-        super().__init__(
-            spec_dim,
-            hidden_dim,
-            filter_len,
-            hop_len,
-            layers,
-            block_layers,
-            kernel_size,
-            is_mask,
-            norm,
-            act,
-        )
-        self.add_spec_results = add_spec_results
-        # refine conv
-        self.refine_conv = nn.ModuleList(
-            [
-                nn.Sequential(
-                    ComplexConvBlock(
-                        spec_dim * 2,
-                        spec_dim * 2,
-                        kernel_size=kernel_size,
-                        padding=kernel_size // 2,
-                        bn_func=self.bn_func,
-                        act_func=self.act_func,
-                    ),
-                    self.bn_func(spec_dim * 2),
-                    self.act_func(),
-                )
-            ]
-            * refine_layers
-        )
-
-    def forward(self, wav):
-        # stft
-        origin_mag, origin_phase = self.log_stft(wav)
-        origin_x = torch.cat([origin_mag, origin_phase], dim=1)
-
-        # prev
-        x = self.prev_conv(origin_x)
-
-        # body
-        # down
-        down_cache = []
-        for idx, block in enumerate(self.down):
-            x = block(x)
-            down_cache.append(x)
-            x = self.down_pool(x)
-
-        # up
-        for idx, block in enumerate(self.up):
-            x = block(x)
-            res = F.interpolate(
-                down_cache[self.layers - (idx + 1)],
-                size=[x.size()[2]],
-                mode="linear",
-                align_corners=False,
-            )
-            x = concat_complex(x, res, dim=1)
-
-        # match spec dimension
-        x = self.out_conv(x)
-        if origin_mag.size(2) != x.size(2):
-            x = F.interpolate(
-                x, size=[origin_mag.size(2)], mode="linear", align_corners=False
-            )
-
-        # refine
-        for idx, refine_module in enumerate(self.refine_conv):
-            x = refine_module(x)
-            mag, phase = x.chunk(2, 1)
-            mag, phase = self.masking(mag, phase, origin_mag, origin_phase)
-            if idx < len(self.refine_conv) - 1:
-                x = torch.cat([mag, phase], dim=1)
-
-        # clamp phase
-        phase = phase.clamp(-np.pi, np.pi)
-
-        out = self.exp_istft(mag, phase)
-        out = self.adjust_diff(out, wav)
-
-        if self.add_spec_results:
-            out = (out, mag, phase)
-
-        return out
-
-
-class _ComplexConvNd(nn.Module):
-    """
-    Implement Complex Convolution
-    A: real weight
-    B: img weight
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride,
-        padding,
-        dilation,
-        transposed,
-        output_padding,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
-        self.dilation = dilation
-        self.output_padding = output_padding
-        self.transposed = transposed
-
-        self.A = self.make_weight(in_channels, out_channels, kernel_size)
-        self.B = self.make_weight(in_channels, out_channels, kernel_size)
-
-        self.reset_parameters()
-
-    def make_weight(self, in_ch, out_ch, kernel_size):
-        if self.transposed:
-            tensor = nn.Parameter(torch.Tensor(in_ch, out_ch // 2, *kernel_size))
-        else:
-            tensor = nn.Parameter(torch.Tensor(out_ch, in_ch // 2, *kernel_size))
-        return tensor
-
-    def reset_parameters(self):
-        # init real weight
-        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.A)
-
-        # init A
-        gain = calculate_gain("leaky_relu", 0)
-        std = gain / np.sqrt(fan_in)
-        bound = np.sqrt(3.0) * std
-
-        with torch.no_grad():
-            # TODO: find more stable initial values
-            self.A.uniform_(-bound * (1 / (np.pi**2)), bound * (1 / (np.pi**2)))
-            #
-            # B is initialized by pi
-            # -pi and pi is too big, so it is powed by -1
-            self.B.uniform_(-1 / np.pi, 1 / np.pi)
-
-
-class ComplexConv1d(_ComplexConvNd):
-    """
-    Complex Convolution 1d
-    """
-
-    def __init__(
-        self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1
-    ):
-        kernel_size = single(kernel_size)
-        stride = single(stride)
-        # edit padding
-        padding = padding
-        dilation = single(dilation)
-        super(ComplexConv1d, self).__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation,
-            False,
-            single(0),
-        )
-
-    def forward(self, x):
-        """
-        Implemented complex convolution using combining 'grouped convolution' and
-        'real / img weight'
-        :param x: data (N, C, T) C is concatenated with C/2 real channels and C/2 idea channels
-        :return: complex conved result
-        """
-        # adopt reflect padding
-        if self.padding:
-            x = F.pad(x, (self.padding, self.padding), "reflect")
-
-        # forward real
-        real_part = F.conv1d(
-            x,
-            self.A,
-            None,
-            stride=self.stride,
-            padding=0,
-            dilation=self.dilation,
-            groups=2,
-        )
-
-        # forward idea
-        spl = self.in_channels // 2
-        weight_B = torch.cat([self.B[:spl].data * (-1), self.B[spl:].data])
-        idea_part = F.conv1d(
-            x,
-            weight_B,
-            None,
-            stride=self.stride,
-            padding=0,
-            dilation=self.dilation,
-            groups=2,
-        )
-
-        return real_part + idea_part
-
-
-class ComplexTransposedConv1d(_ComplexConvNd):
-    """
-    Complex Transposed Convolution 1d
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        output_padding=0,
-        dilation=1,
-    ):
-        kernel_size = single(kernel_size)
-        stride = single(stride)
-        padding = padding
-        dilation = single(dilation)
-        super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation,
-            True,
-            output_padding,
-        )
-
-    def forward(self, x, output_size=None):
-        """
-        Implemented complex transposed convolution using combining 'grouped convolution'
-        and 'real / img weight'
-        :param x: data (N, C, T) C is concatenated with C/2 real channels and C/2 idea channels
-        :return: complex transposed convolution result
-        """
-        # forward real
-        if self.padding:
-            x = F.pad(x, (self.padding, self.padding), "reflect")
-
-        real_part = F.conv_transpose1d(
-            x,
-            self.A,
-            None,
-            stride=self.stride,
-            padding=0,
-            dilation=self.dilation,
-            groups=2,
-        )
-
-        # forward idea
-        spl = self.out_channels // 2
-        weight_B = torch.cat([self.B[:spl] * (-1), self.B[spl:]])
-        idea_part = F.conv_transpose1d(
-            x,
-            weight_B,
-            None,
-            stride=self.stride,
-            padding=0,
-            dilation=self.dilation,
-            groups=2,
-        )
-
-        if self.output_padding:
-            real_part = F.pad(
-                real_part, (self.output_padding, self.output_padding), "reflect"
-            )
-            idea_part = F.pad(
-                idea_part, (self.output_padding, self.output_padding), "reflect"
-            )
-
-        return real_part + idea_part
-
-
-class ComplexActLayer(nn.Module):
-    """
-    Activation differently 'real' part and 'img' part
-    In implemented DCUnet on this repository, Real part is activated to log space.
-    And Phase(img) part, it is distributed in [-pi, pi]...
-    """
-
-    def forward(self, x):
-        real, img = x.chunk(2, 1)
-        return torch.cat([F.leaky_relu(real), torch.tanh(img) * np.pi], dim=1)
-
-
-class STFT(nn.Module):
-    """
-    Re-construct stft for calculating backward operation
-    refer on : https://github.com/pseeth/torch-stft/blob/master/torch_stft/stft.py
-    """
-
-    def __init__(
-        self,
-        filter_length: int = 1024,
-        hop_length: int = 512,
-        win_length: int = None,
-        window: str = "hann",
-    ):
-        super().__init__()
-        self.filter_length = filter_length
-        self.hop_length = hop_length
-        self.win_length = win_length if win_length else filter_length
-        self.window = window
-        self.pad_amount = self.filter_length // 2
-
-        # make fft window
-        assert filter_length >= self.win_length
-        # get window and zero center pad it to filter_length
-        fft_window = get_window(window, self.win_length, fftbins=True)
-        fft_window = pad_center(fft_window, filter_length)
-        fft_window = torch.from_numpy(fft_window).float()
-
-        # calculate fourer_basis
-        cut_off = int((self.filter_length / 2 + 1))
-        fourier_basis = np.fft.fft(np.eye(self.filter_length))
-        fourier_basis = np.vstack(
-            [np.real(fourier_basis[:cut_off, :]), np.imag(fourier_basis[:cut_off, :])]
-        )
-
-        # make forward & inverse basis
-        self.register_buffer("square_window", fft_window**2)
-
-        forward_basis = torch.FloatTensor(fourier_basis[:, np.newaxis, :]) * fft_window
-        inverse_basis = (
-            torch.FloatTensor(
-                np.linalg.pinv(self.filter_length / self.hop_length * fourier_basis).T[
-                    :, np.newaxis, :
-                ]
-            )
-            * fft_window
-        )
-        # torch.pinverse has a bug, so at this time, it is separated into two parts..
-        self.register_buffer("forward_basis", forward_basis)
-        self.register_buffer("inverse_basis", inverse_basis)
-
-    def transform(self, wav: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        # reflect padding
-        wav = wav.unsqueeze(1).unsqueeze(1)
-        wav = F.pad(
-            wav, (self.pad_amount, self.pad_amount, 0, 0), mode="reflect"
-        ).squeeze(1)
-
-        # conv
-        forward_trans = F.conv1d(
-            wav, self.forward_basis, stride=self.hop_length, padding=0
-        )
-        real_part, imag_part = forward_trans.chunk(2, 1)
-
-        return torch.sqrt(real_part**2 + imag_part**2), torch.atan2(
-            imag_part.data, real_part.data
-        )
-
-    def inverse(
-        self, magnitude: torch.Tensor, phase: torch.Tensor, eps: float = 1e-9
-    ) -> torch.Tensor:
-        comp = torch.cat(
-            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
-        )
-        inverse_transform = F.conv_transpose1d(
-            comp, self.inverse_basis, stride=self.hop_length, padding=0
-        )
-
-        # remove window effect
-        n_frames = comp.size(-1)
-        inverse_size = inverse_transform.size(-1)
-
-        window_filter = torch.ones(1, 1, n_frames).type_as(inverse_transform)
-
-        weight = self.square_window[: self.filter_length].unsqueeze(0).unsqueeze(0)
-        window_filter = F.conv_transpose1d(
-            window_filter, weight, stride=self.hop_length, padding=0
-        )
-        window_filter = window_filter.squeeze()[:inverse_size] + eps
-
-        inverse_transform /= window_filter
-
-        # scale by hop ratio
-        inverse_transform *= self.filter_length / self.hop_length
-
-        return inverse_transform[..., self.pad_amount : -self.pad_amount].squeeze(1)

remfx/dptnet.py

@@ -1,459 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.nn.modules.container import ModuleList
-from torch.nn.modules.activation import MultiheadAttention
-from torch.nn.modules.dropout import Dropout
-from torch.nn.modules.linear import Linear
-from torch.nn.modules.rnn import LSTM
-from torch.nn.modules.normalization import LayerNorm
-from torch.autograd import Variable
-import copy
-import math
-
-
-# adapted from https://github.com/ujscjj/DPTNet
-
-
-class DPTNet_base(nn.Module):
-    def __init__(
-        self,
-        enc_dim,
-        feature_dim,
-        hidden_dim,
-        layer,
-        segment_size=250,
-        nspk=2,
-        win_len=2,
-    ):
-        super().__init__()
-        # parameters
-        self.window = win_len
-        self.stride = self.window // 2
-
-        self.enc_dim = enc_dim
-        self.feature_dim = feature_dim
-        self.hidden_dim = hidden_dim
-        self.segment_size = segment_size
-
-        self.layer = layer
-        self.num_spk = nspk
-        self.eps = 1e-8
-
-        self.dpt_encoder = DPTEncoder(
-            n_filters=enc_dim,
-            window_size=win_len,
-        )
-        self.enc_LN = nn.GroupNorm(1, self.enc_dim, eps=1e-8)
-        self.dpt_separation = DPTSeparation(
-            self.enc_dim,
-            self.feature_dim,
-            self.hidden_dim,
-            self.num_spk,
-            self.layer,
-            self.segment_size,
-        )
-
-        self.mask_conv1x1 = nn.Conv1d(self.feature_dim, self.enc_dim, 1, bias=False)
-        self.decoder = DPTDecoder(n_filters=enc_dim, window_size=win_len)
-
-    def forward(self, mix):
-        """
-        mix: shape (batch, T)
-        """
-        batch_size = mix.shape[0]
-        mix = self.dpt_encoder(mix)  # (B, E, L)
-
-        score_ = self.enc_LN(mix)  # B, E, L
-        score_ = self.dpt_separation(score_)  # B, nspk, T, N
-        score_ = (
-            score_.view(batch_size * self.num_spk, -1, self.feature_dim)
-            .transpose(1, 2)
-            .contiguous()
-        )  # B*nspk, N, T
-        score = self.mask_conv1x1(score_)  # [B*nspk, N, L] -> [B*nspk, E, L]
-        score = score.view(
-            batch_size, self.num_spk, self.enc_dim, -1
-        )  # [B*nspk, E, L] -> [B, nspk, E, L]
-        est_mask = F.relu(score)
-
-        est_source = self.decoder(
-            mix, est_mask
-        )  # [B, E, L] + [B, nspk, E, L]--> [B, nspk, T]
-
-        return est_source
-
-
-class DPTEncoder(nn.Module):
-    def __init__(self, n_filters: int = 64, window_size: int = 2):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            1, n_filters, kernel_size=window_size, stride=window_size // 2, bias=False
-        )
-
-    def forward(self, x):
-        x = x.unsqueeze(1)
-        x = F.relu(self.conv(x))
-        return x
-
-
-class TransformerEncoderLayer(torch.nn.Module):
-    def __init__(
-        self, d_model, nhead, hidden_size, dim_feedforward, dropout, activation="relu"
-    ):
-        super(TransformerEncoderLayer, self).__init__()
-        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
-
-        # Implementation of improved part
-        self.lstm = LSTM(d_model, hidden_size, 1, bidirectional=True)
-        self.dropout = Dropout(dropout)
-        self.linear = Linear(hidden_size * 2, d_model)
-
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.dropout1 = Dropout(dropout)
-        self.dropout2 = Dropout(dropout)
-
-        self.activation = _get_activation_fn(activation)
-
-    def __setstate__(self, state):
-        if "activation" not in state:
-            state["activation"] = F.relu
-        super(TransformerEncoderLayer, self).__setstate__(state)
-
-    def forward(self, src, src_mask=None, src_key_padding_mask=None):
-        r"""Pass the input through the encoder layer.
-        Args:
-            src: the sequnce to the encoder layer (required).
-            src_mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-        Shape:
-            see the docs in Transformer class.
-        """
-        src2 = self.self_attn(
-            src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
-        )[0]
-        src = src + self.dropout1(src2)
-        src = self.norm1(src)
-        src2 = self.linear(self.dropout(self.activation(self.lstm(src)[0])))
-        src = src + self.dropout2(src2)
-        src = self.norm2(src)
-        return src
-
-
-def _get_clones(module, N):
-    return ModuleList([copy.deepcopy(module) for i in range(N)])
-
-
-def _get_activation_fn(activation):
-    if activation == "relu":
-        return F.relu
-    elif activation == "gelu":
-        return F.gelu
-
-    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
-
-
-class SingleTransformer(nn.Module):
-    """
-    Container module for a single Transformer layer.
-    args: input_size: int, dimension of the input feature.
-    The input should have shape (batch, seq_len, input_size).
-    """
-
-    def __init__(self, input_size, hidden_size, dropout):
-        super(SingleTransformer, self).__init__()
-        self.transformer = TransformerEncoderLayer(
-            d_model=input_size,
-            nhead=4,
-            hidden_size=hidden_size,
-            dim_feedforward=hidden_size * 2,
-            dropout=dropout,
-        )
-
-    def forward(self, input):
-        # input shape: batch, seq, dim
-        output = input
-        transformer_output = (
-            self.transformer(output.permute(1, 0, 2).contiguous())
-            .permute(1, 0, 2)
-            .contiguous()
-        )
-        return transformer_output
-
-
-# dual-path transformer
-class DPT(nn.Module):
-    """
-    Deep dual-path transformer.
-    args:
-        input_size: int, dimension of the input feature. The input should have shape
-                    (batch, seq_len, input_size).
-        hidden_size: int, dimension of the hidden state.
-        output_size: int, dimension of the output size.
-        num_layers: int, number of stacked Transformer layers. Default is 1.
-        dropout: float, dropout ratio. Default is 0.
-    """
-
-    def __init__(self, input_size, hidden_size, output_size, num_layers=1, dropout=0):
-        super(DPT, self).__init__()
-
-        self.input_size = input_size
-        self.output_size = output_size
-        self.hidden_size = hidden_size
-
-        # dual-path transformer
-        self.row_transformer = nn.ModuleList([])
-        self.col_transformer = nn.ModuleList([])
-        for i in range(num_layers):
-            self.row_transformer.append(
-                SingleTransformer(input_size, hidden_size, dropout)
-            )
-            self.col_transformer.append(
-                SingleTransformer(input_size, hidden_size, dropout)
-            )
-
-        # output layer
-        self.output = nn.Sequential(nn.PReLU(), nn.Conv2d(input_size, output_size, 1))
-
-    def forward(self, input):
-        # input shape: batch, N, dim1, dim2
-        # apply transformer on dim1 first and then dim2
-        # output shape: B, output_size, dim1, dim2
-        # input = input.to(device)
-        batch_size, _, dim1, dim2 = input.shape
-        output = input
-        for i in range(len(self.row_transformer)):
-            row_input = (
-                output.permute(0, 3, 2, 1)
-                .contiguous()
-                .view(batch_size * dim2, dim1, -1)
-            )  # B*dim2, dim1, N
-            row_output = self.row_transformer[i](row_input)  # B*dim2, dim1, H
-            row_output = (
-                row_output.view(batch_size, dim2, dim1, -1)
-                .permute(0, 3, 2, 1)
-                .contiguous()
-            )  # B, N, dim1, dim2
-            output = row_output
-
-            col_input = (
-                output.permute(0, 2, 3, 1)
-                .contiguous()
-                .view(batch_size * dim1, dim2, -1)
-            )  # B*dim1, dim2, N
-            col_output = self.col_transformer[i](col_input)  # B*dim1, dim2, H
-            col_output = (
-                col_output.view(batch_size, dim1, dim2, -1)
-                .permute(0, 3, 1, 2)
-                .contiguous()
-            )  # B, N, dim1, dim2
-            output = col_output
-
-        output = self.output(output)  # B, output_size, dim1, dim2
-
-        return output
-
-
-# base module for deep DPT
-class DPT_base(nn.Module):
-    def __init__(
-        self, input_dim, feature_dim, hidden_dim, num_spk=2, layer=6, segment_size=250
-    ):
-        super(DPT_base, self).__init__()
-
-        self.input_dim = input_dim
-        self.feature_dim = feature_dim
-        self.hidden_dim = hidden_dim
-
-        self.layer = layer
-        self.segment_size = segment_size
-        self.num_spk = num_spk
-
-        self.eps = 1e-8
-
-        # bottleneck
-        self.BN = nn.Conv1d(self.input_dim, self.feature_dim, 1, bias=False)
-
-        # DPT model
-        self.DPT = DPT(
-            self.feature_dim,
-            self.hidden_dim,
-            self.feature_dim * self.num_spk,
-            num_layers=layer,
-        )
-
-    def pad_segment(self, input, segment_size):
-        # input is the features: (B, N, T)
-        batch_size, dim, seq_len = input.shape
-        segment_stride = segment_size // 2
-
-        rest = segment_size - (segment_stride + seq_len % segment_size) % segment_size
-        if rest > 0:
-            pad = Variable(torch.zeros(batch_size, dim, rest)).type(input.type())
-            input = torch.cat([input, pad], 2)
-
-        pad_aux = Variable(torch.zeros(batch_size, dim, segment_stride)).type(
-            input.type()
-        )
-        input = torch.cat([pad_aux, input, pad_aux], 2)
-
-        return input, rest
-
-    def split_feature(self, input, segment_size):
-        # split the feature into chunks of segment size
-        # input is the features: (B, N, T)
-
-        input, rest = self.pad_segment(input, segment_size)
-        batch_size, dim, seq_len = input.shape
-        segment_stride = segment_size // 2
-
-        segments1 = (
-            input[:, :, :-segment_stride]
-            .contiguous()
-            .view(batch_size, dim, -1, segment_size)
-        )
-        segments2 = (
-            input[:, :, segment_stride:]
-            .contiguous()
-            .view(batch_size, dim, -1, segment_size)
-        )
-        segments = (
-            torch.cat([segments1, segments2], 3)
-            .view(batch_size, dim, -1, segment_size)
-            .transpose(2, 3)
-        )
-
-        return segments.contiguous(), rest
-
-    def merge_feature(self, input, rest):
-        # merge the splitted features into full utterance
-        # input is the features: (B, N, L, K)
-
-        batch_size, dim, segment_size, _ = input.shape
-        segment_stride = segment_size // 2
-        input = (
-            input.transpose(2, 3)
-            .contiguous()
-            .view(batch_size, dim, -1, segment_size * 2)
-        )  # B, N, K, L
-
-        input1 = (
-            input[:, :, :, :segment_size]
-            .contiguous()
-            .view(batch_size, dim, -1)[:, :, segment_stride:]
-        )
-        input2 = (
-            input[:, :, :, segment_size:]
-            .contiguous()
-            .view(batch_size, dim, -1)[:, :, :-segment_stride]
-        )
-
-        output = input1 + input2
-        if rest > 0:
-            output = output[:, :, :-rest]
-
-        return output.contiguous()  # B, N, T
-
-    def forward(self, input):
-        pass
-
-
-class DPTSeparation(DPT_base):
-    def __init__(self, *args, **kwargs):
-        super(DPTSeparation, self).__init__(*args, **kwargs)
-
-        # gated output layer
-        self.output = nn.Sequential(
-            nn.Conv1d(self.feature_dim, self.feature_dim, 1), nn.Tanh()
-        )
-        self.output_gate = nn.Sequential(
-            nn.Conv1d(self.feature_dim, self.feature_dim, 1), nn.Sigmoid()
-        )
-
-    def forward(self, input):
-        # input = input.to(device)
-        # input: (B, E, T)
-        batch_size, E, seq_length = input.shape
-
-        enc_feature = self.BN(input)  # (B, E, L)-->(B, N, L)
-        # split the encoder output into overlapped, longer segments
-        enc_segments, enc_rest = self.split_feature(
-            enc_feature, self.segment_size
-        )  # B, N, L, K: L is the segment_size
-        # print('enc_segments.shape {}'.format(enc_segments.shape))
-        # pass to DPT
-        output = self.DPT(enc_segments).view(
-            batch_size * self.num_spk, self.feature_dim, self.segment_size, -1
-        )  # B*nspk, N, L, K
-
-        # overlap-and-add of the outputs
-        output = self.merge_feature(output, enc_rest)  # B*nspk, N, T
-
-        # gated output layer for filter generation
-        bf_filter = self.output(output) * self.output_gate(output)  # B*nspk, K, T
-        bf_filter = (
-            bf_filter.transpose(1, 2)
-            .contiguous()
-            .view(batch_size, self.num_spk, -1, self.feature_dim)
-        )  # B, nspk, T, N
-
-        return bf_filter
-
-
-class DPTDecoder(nn.Module):
-    def __init__(self, n_filters: int = 64, window_size: int = 2):
-        super().__init__()
-        self.W = window_size
-        self.basis_signals = nn.Linear(n_filters, window_size, bias=False)
-
-    def forward(self, mixture, mask):
-        """
-        mixture: (batch, n_filters, L)
-        mask: (batch, sources, n_filters, L)
-        """
-        source_w = torch.unsqueeze(mixture, 1) * mask  # [B, C, E, L]
-        source_w = torch.transpose(source_w, 2, 3)  # [B, C, L, E]
-        # S = DV
-        est_source = self.basis_signals(source_w)  # [B, C, L, W]
-        est_source = overlap_and_add(est_source, self.W // 2)  # B x C x T
-        return est_source
-
-
-def overlap_and_add(signal, frame_step):
-    """Reconstructs a signal from a framed representation.
-    Adds potentially overlapping frames of a signal with shape
-    `[..., frames, frame_length]`, offsetting subsequent frames by `frame_step`.
-    The resulting tensor has shape `[..., output_size]` where
-        output_size = (frames - 1) * frame_step + frame_length
-    Args:
-        signal: A [..., frames, frame_length] Tensor.
-        All dimensions may be unknown, and rank must be at least 2.
-        frame_step: An integer denoting overlap offsets. Must be less than or equal to frame_length.
-    Returns:
-        A Tensor with shape [..., output_size] containing the overlap-added frames of signal's
-        inner-most two dimensions.
-        output_size = (frames - 1) * frame_step + frame_length
-    Based on https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/signal/python/ops/reconstruction_ops.py
-    """
-    outer_dimensions = signal.size()[:-2]
-    frames, frame_length = signal.size()[-2:]
-
-    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
-    subframe_step = frame_step // subframe_length
-    subframes_per_frame = frame_length // subframe_length
-    output_size = frame_step * (frames - 1) + frame_length
-    output_subframes = output_size // subframe_length
-
-    subframe_signal = signal.reshape(*outer_dimensions, -1, subframe_length)
-
-    frame = torch.arange(0, output_subframes).unfold(
-        0, subframes_per_frame, subframe_step
-    )
-    frame = signal.new_tensor(frame).long()  # signal may in GPU or CPU
-    frame = frame.contiguous().view(-1)
-
-    result = signal.new_zeros(*outer_dimensions, output_subframes, subframe_length)
-    result.index_add_(-2, frame, subframe_signal)
-    result = result.view(*outer_dimensions, -1)
-    return result

remfx/models.py

@@ -226,7 +226,7 @@ class DCUNetModel(nn.Module):
 
     def forward(self, batch):
         x, target = batch
-        output = self.model(x.squeeze(1))  # B x 1 x T
+        output = self.model(x.squeeze(1))  # B x T
         # Crop target to match output
         if output.shape[-1] < target.shape[-1]:
             target = causal_crop(target, output.shape[-1])
@@ -234,7 +234,7 @@ class DCUNetModel(nn.Module):
         return loss, output
 
     def sample(self, x: Tensor) -> Tensor:
-        output = self.model(x.squeeze(1))  # B x 1 x T
+        output = self.model(x.squeeze(1))  # B x T
         return output