Merge pull request #37 from mhrice/dsd100-dataset

cfg/config.yaml

@@ -53,9 +53,10 @@ callbacks:
     _target_: remfx.callbacks.MetricCallback
 
 datamodule:
-  _target_: remfx.datasets.VocalSetDatamodule
+  _target_: remfx.datasets.EffectDatamodule
   train_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
+    total_chunks: 8000
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -70,7 +71,8 @@ datamodule:
     render_files: ${render_files}
     render_root: ${render_root}
   val_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
+    total_chunks: 1000
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -85,7 +87,8 @@ datamodule:
     render_files: ${render_files}
     render_root: ${render_root}
   test_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
+    total_chunks: 1000
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}

cfg/effects/all.yaml

@@ -4,13 +4,19 @@ effects:
   chorus:
       _target_: remfx.effects.RandomPedalboardChorus
       sample_rate: ${sample_rate}
+      min_rate_hz: 0.25
+      max_rate_hz: 1.5
+      min_feedback: 0.1
+      max_feedback: 0.4
       min_depth: 0.2
-      min_mix: 0.3
+      max_depth: 0.6
+      min_mix: 0.15
+      max_mix: 0.4
   distortion:
       _target_: remfx.effects.RandomPedalboardDistortion
       sample_rate: ${sample_rate}
-      min_drive_db: 10
-      max_drive_db: 50
+      min_drive_db: 7
+      max_drive_db: 25
   compressor:
       _target_: remfx.effects.RandomPedalboardCompressor
       sample_rate: ${sample_rate}
@@ -26,7 +32,7 @@ effects:
       min_damping: 0.2
       max_damping: 1.0
       min_wet_dry: 0.2
-      max_wet_dry: 0.8
+      max_wet_dry: 0.6
       min_width: 0.2
       max_width: 1.0
   delay:
@@ -35,6 +41,6 @@ effects:
       min_delay_seconds: 0.1
       max_delay_sconds: 1.0
       min_feedback: 0.05
-      max_feedback: 0.6
-      min_mix: 0.2
-      max_mix: 0.7
+      max_feedback: 0.3
+      min_mix: 0.1
+      max_mix: 0.35

cfg/exp/default.yaml

@@ -1,6 +1,6 @@
 # @package _global_
 defaults:
-  - override /model: demucs
+  - override /model: umx
   - override /effects: all
 seed: 12345
 sample_rate: 48000

cfg/exp/dist.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+  - override /model: umx
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+render_files: True
+render_root: "/scratch/EffectSet"
+accelerator: "gpu"
+log_audio: True
+# Effects
+max_kept_effects: 5
+max_removed_effects: -1
+shuffle_kept_effects: True
+shuffle_removed_effects: False
+num_classes: 5
+effects_to_use:
+  - compressor
+  - distortion
+  - reverb
+  - chorus
+  - delay
+effects_to_remove:
+  - distortion
+datamodule:
+  batch_size: 16
+  num_workers: 8

cfg/model/dcunet.yaml

@@ -9,16 +9,8 @@ model:
   sample_rate: ${sample_rate}
   network:
     _target_: remfx.models.DCUNetModel
-    spec_dim: 257
-    hidden_dim: 768
-    filter_len: 512
-    hop_len: 64
-    block_layers: 4
-    layers: 4
-    kernel_size: 3
-    refine_layers: 1
-    is_mask: True
-    norm: 'ins'
-    act: 'comp'
+    architecture: "DCUNet-10"
+    stft_kernel_size: 512
+    fix_length_mode: "pad"
     sample_rate: ${sample_rate}
     num_bins: 1025

cfg/model/dptnet.yaml

@@ -9,12 +9,14 @@ model:
   sample_rate: ${sample_rate}
   network:
     _target_: remfx.models.DPTNetModel
-    enc_dim: 256
-    feature_dim: 64
-    hidden_dim: 128
-    layer: 6
-    segment_size: 250
-    nspk: 1
-    win_len: 2
+    n_src: 1
+    in_chan: 64
+    out_chan: 64
+    chunk_size: 100
+    n_repeats: 2
+    fb_name: "free"
+    kernel_size: 16
+    n_filters: 64
+    stride: 8
     sample_rate: ${sample_rate}
     num_bins: 1025

cfg/model/tcn.yaml

@@ -11,12 +11,12 @@ model:
     _target_: remfx.models.TCNModel
     ninputs: 1
     noutputs: 1
-    nblocks: 4
+    nblocks: 20
     channel_growth: 0
-    channel_width: 32
-    kernel_size: 13
+    channel_width: 64
+    kernel_size: 7
     stack_size: 10
-    dilation_growth: 10
+    dilation_growth: 2
     condition: False
     latent_dim: 2
     norm_type: "identity"

remfx/datasets.py

@@ -5,20 +5,20 @@ import torch
 import shutil
 import torchaudio
 import pytorch_lightning as pl
-
+import random
 from tqdm import tqdm
 from pathlib import Path
 from remfx import effects
 from ordered_set import OrderedSet
 from typing import Any, List, Dict
 from torch.utils.data import Dataset, DataLoader
-from remfx.utils import create_sequential_chunks
+from remfx.utils import select_random_chunk
 
 
 # https://zenodo.org/record/1193957 -> VocalSet
 
 ALL_EFFECTS = effects.Pedalboard_Effects
-print(ALL_EFFECTS)
+# print(ALL_EFFECTS)
 
 
 vocalset_splits = {
@@ -55,6 +55,11 @@ idmt_bass_splits = {
     "val": ["VIF"],
     "test": ["VIS"],
 }
+dsd_100_splits = {
+    "train": ["train"],
+    "val": ["val"],
+    "test": ["test"],
+}
 idmt_drums_splits = {
     "train": ["WaveDrum02", "TechnoDrum01"],
     "val": ["RealDrum01"],
@@ -76,7 +81,7 @@ def locate_files(root: str, mode: str):
         for singer_dir in singer_dirs:
             files += glob.glob(os.path.join(singer_dir, "**", "**", "*.wav"))
         print(f"Found {len(files)} files in VocalSet {mode}.")
-        file_list += sorted(files)
+        file_list.append(sorted(files))
     # ------------------------- GuitarSet -------------------------
     guitarset_dir = os.path.join(root, "audio_mono-mic")
     if os.path.isdir(guitarset_dir):
@@ -87,37 +92,46 @@ def locate_files(root: str, mode: str):
             if os.path.basename(f).split("_")[0] in guitarset_splits[mode]
         ]
         print(f"Found {len(files)} files in GuitarSet {mode}.")
-        file_list += sorted(files)
-    # ------------------------- IDMT-SMT-GUITAR -------------------------
-    idmt_smt_guitar_dir = os.path.join(root, "IDMT-SMT-GUITAR_V2")
-    if os.path.isdir(idmt_smt_guitar_dir):
-        files = glob.glob(
-            os.path.join(
-                idmt_smt_guitar_dir, "IDMT-SMT-GUITAR_V2", "dataset4", "**", "*.wav"
-            ),
-            recursive=True,
-        )
-        files = [
-            f
-            for f in files
-            if os.path.basename(f).split("_")[0] in idmt_guitar_splits[mode]
-        ]
-        file_list += sorted(files)
-        print(f"Found {len(files)} files in IDMT-SMT-Guitar {mode}.")
+        file_list.append(sorted(files))
+    # # ------------------------- IDMT-SMT-GUITAR -------------------------
+    # idmt_smt_guitar_dir = os.path.join(root, "IDMT-SMT-GUITAR_V2")
+    # if os.path.isdir(idmt_smt_guitar_dir):
+    #     files = glob.glob(
+    #         os.path.join(
+    #             idmt_smt_guitar_dir, "IDMT-SMT-GUITAR_V2", "dataset4", "**", "*.wav"
+    #         ),
+    #         recursive=True,
+    #     )
+    #     files = [
+    #         f
+    #         for f in files
+    #         if os.path.basename(f).split("_")[0] in idmt_guitar_splits[mode]
+    #     ]
+    #     file_list.append(sorted(files))
+    #     print(f"Found {len(files)} files in IDMT-SMT-Guitar {mode}.")
     # ------------------------- IDMT-SMT-BASS -------------------------
-    idmt_smt_bass_dir = os.path.join(root, "IDMT-SMT-BASS")
-    if os.path.isdir(idmt_smt_bass_dir):
+    # idmt_smt_bass_dir = os.path.join(root, "IDMT-SMT-BASS")
+    # if os.path.isdir(idmt_smt_bass_dir):
+    #     files = glob.glob(
+    #         os.path.join(idmt_smt_bass_dir, "**", "*.wav"),
+    #         recursive=True,
+    #     )
+    #     files = [
+    #         f
+    #         for f in files
+    #         if os.path.basename(os.path.dirname(f)) in idmt_bass_splits[mode]
+    #     ]
+    #     file_list.append(sorted(files))
+    #     print(f"Found {len(files)} files in IDMT-SMT-Bass {mode}.")
+    # ------------------------- DSD100 ---------------------------------
+    dsd_100_dir = os.path.join(root, "DSD100")
+    if os.path.isdir(dsd_100_dir):
         files = glob.glob(
-            os.path.join(idmt_smt_bass_dir, "**", "*.wav"),
+            os.path.join(dsd_100_dir, mode, "**", "*.wav"),
             recursive=True,
         )
-        files = [
-            f
-            for f in files
-            if os.path.basename(os.path.dirname(f)) in idmt_bass_splits[mode]
-        ]
-        file_list += sorted(files)
-        print(f"Found {len(files)} files in IDMT-SMT-Bass {mode}.")
+        file_list.append(sorted(files))
+        print(f"Found {len(files)} files in DSD100 {mode}.")
     # ------------------------- IDMT-SMT-DRUMS -------------------------
     idmt_smt_drums_dir = os.path.join(root, "IDMT-SMT-DRUMS-V2")
     if os.path.isdir(idmt_smt_drums_dir):
@@ -127,18 +141,19 @@ def locate_files(root: str, mode: str):
             for f in files
             if os.path.basename(f).split("_")[0] in idmt_drums_splits[mode]
         ]
-        file_list += sorted(files)
+        file_list.append(sorted(files))
         print(f"Found {len(files)} files in IDMT-SMT-Drums {mode}.")
 
     return file_list
 
 
-class VocalSet(Dataset):
+class EffectDataset(Dataset):
     def __init__(
         self,
         root: str,
         sample_rate: int,
         chunk_size: int = 262144,
+        total_chunks: int = 1000,
         effect_modules: List[Dict[str, torch.nn.Module]] = None,
         effects_to_use: List[str] = None,
         effects_to_remove: List[str] = None,
@@ -156,6 +171,7 @@ class VocalSet(Dataset):
         self.root = Path(root)
         self.render_root = Path(render_root)
         self.chunk_size = chunk_size
+        self.total_chunks = total_chunks
         self.sample_rate = sample_rate
         self.mode = mode
         self.max_kept_effects = max_kept_effects
@@ -184,42 +200,40 @@ class VocalSet(Dataset):
                     sys.exit()
                 shutil.rmtree(self.proc_root)
 
-        self.num_chunks = 0
-        print("Total files:", len(self.files))
+        print("Total datasets:", len(self.files))
         print("Processing files...")
         if render_files:
             # Split audio file into chunks, resample, then apply random effects
             self.proc_root.mkdir(parents=True, exist_ok=True)
-            for audio_file in tqdm(self.files, total=len(self.files)):
-                chunks, orig_sr = create_sequential_chunks(audio_file, self.chunk_size)
-                for chunk in chunks:
-                    resampled_chunk = torchaudio.functional.resample(
-                        chunk, orig_sr, sample_rate
+            for num_chunk in tqdm(range(self.total_chunks)):
+                chunk = None
+                random_dataset_choice = random.choice(self.files)
+                while chunk is None:
+                    random_file_choice = random.choice(random_dataset_choice)
+                    chunk = select_random_chunk(
+                        random_file_choice, self.chunk_size, self.sample_rate
                     )
-                    if resampled_chunk.shape[-1] < chunk_size:
-                        # Skip if chunk is too small
-                        continue
 
-                    dry, wet, dry_effects, wet_effects = self.process_effects(
-                        resampled_chunk
-                    )
-                    output_dir = self.proc_root / str(self.num_chunks)
-                    output_dir.mkdir(exist_ok=True)
-                    torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
-                    torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
-                    torch.save(dry_effects, output_dir / "dry_effects.pt")
-                    torch.save(wet_effects, output_dir / "wet_effects.pt")
-                    self.num_chunks += 1
+                # Sum to mono
+                if chunk.shape[0] > 1:
+                    chunk = chunk.sum(0, keepdim=True)
+
+                dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
+                output_dir = self.proc_root / str(num_chunk)
+                output_dir.mkdir(exist_ok=True)
+                torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
+                torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
+                torch.save(dry_effects, output_dir / "dry_effects.pt")
+                torch.save(wet_effects, output_dir / "wet_effects.pt")
+
+            print("Finished rendering")
         else:
-            self.num_chunks = len(list(self.proc_root.iterdir()))
+            self.total_chunks = len(list(self.proc_root.iterdir()))
 
-        print(
-            f"Found {len(self.files)} {self.mode} files .\n"
-            f"Total chunks: {self.num_chunks}"
-        )
+        print("Total chunks:", self.total_chunks)
 
     def __len__(self):
-        return self.num_chunks
+        return self.total_chunks
 
     def __getitem__(self, idx):
         input_file = self.proc_root / str(idx) / "input.wav"
@@ -281,7 +295,7 @@ class VocalSet(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]
@@ -292,7 +306,8 @@ class VocalSet(Dataset):
         # Apply
         dry_labels = []
         for effect in effects_to_apply:
-            dry = effect(dry)
+            # Normalize in-between effects
+            dry = self.normalize(effect(dry))
             dry_labels.append(ALL_EFFECTS.index(type(effect)))
 
         # Apply effects_to_remove
@@ -315,7 +330,8 @@ class VocalSet(Dataset):
 
         wet_labels = []
         for effect in effects_to_apply:
-            wet = effect(wet)
+            # Normalize in-between effects
+            wet = self.normalize(effect(wet))
             wet_labels.append(ALL_EFFECTS.index(type(effect)))
 
         wet_labels_tensor = torch.zeros(len(ALL_EFFECTS))
@@ -334,7 +350,7 @@ class VocalSet(Dataset):
         return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
 
 
-class VocalSetDatamodule(pl.LightningDataModule):
+class EffectDatamodule(pl.LightningDataModule):
     def __init__(
         self,
         train_dataset,

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

@@ -2,7 +2,6 @@ import torch
 import torchmetrics
 import pytorch_lightning as pl
 from torch import Tensor, nn
-from torch.nn import functional as F
 from torchaudio.models import HDemucs
 from audio_diffusion_pytorch import DiffusionModel
 from auraloss.time import SISDRLoss
@@ -14,6 +13,7 @@ from remfx.dptnet import DPTNet_base
 from remfx.dcunet import RefineSpectrogramUnet
 from remfx.tcn import TCN
 from remfx.utils import causal_crop
+import asteroid
 
 
 class RemFX(pl.LightningModule):
@@ -85,6 +85,9 @@ class RemFX(pl.LightningModule):
         x, y, _, _ = batch  # x, y = (B, C, T), (B, C, T)
 
         loss, output = self.model((x, y))
+        # Crop target to match output
+        if output.shape[-1] < y.shape[-1]:
+            y = causal_crop(y, output.shape[-1])
         self.log(f"{mode}_loss", loss)
         # Metric logging
         with torch.no_grad():
@@ -195,7 +198,7 @@ class DiffusionGenerationModel(nn.Module):
 class DPTNetModel(nn.Module):
     def __init__(self, sample_rate, num_bins, **kwargs):
         super().__init__()
-        self.model = DPTNet_base(**kwargs)
+        self.model = asteroid.models.dptnet.DPTNet(**kwargs)
         self.num_bins = num_bins
         self.mrstftloss = MultiResolutionSTFTLoss(
             n_bins=self.num_bins, sample_rate=sample_rate
@@ -215,7 +218,7 @@ class DPTNetModel(nn.Module):
 class DCUNetModel(nn.Module):
     def __init__(self, sample_rate, num_bins, **kwargs):
         super().__init__()
-        self.model = RefineSpectrogramUnet(**kwargs)
+        self.model = asteroid.models.DCUNet(**kwargs)
         self.mrstftloss = MultiResolutionSTFTLoss(
             n_bins=num_bins, sample_rate=sample_rate
         )
@@ -223,7 +226,7 @@ class DCUNetModel(nn.Module):
 
     def forward(self, batch):
         x, target = batch
-        output = self.model(x.squeeze(1)).unsqueeze(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])
@@ -231,7 +234,7 @@ class DCUNetModel(nn.Module):
         return loss, output
 
     def sample(self, x: Tensor) -> Tensor:
-        output = self.model(x.squeeze(1)).unsqueeze(1)  # B x 1 x T
+        output = self.model(x.squeeze(1))  # B x T
         return output
 
 

remfx/tcn.py

@@ -128,10 +128,7 @@ class TCN(nn.Module):
         x_in = x
         for _, block in enumerate(self.process_blocks):
             x = block(x)
-        # y_hat = torch.tanh(self.output(x))
-        x_in = causal_crop(x_in, x.shape[-1])
-        gain_ln = self.output(x)
-        y_hat = torch.tanh(gain_ln * x_in)
+        y_hat = torch.tanh(self.output(x))
         return y_hat
 
     def compute_receptive_field(self):

remfx/utils.py

@@ -127,10 +127,10 @@ def create_random_chunks(
 
 
 def create_sequential_chunks(
-    audio_file: str, chunk_size: int
-) -> Tuple[List[Tuple[int, int]], int]:
-    """Create sequential chunks of size chunk_size (seconds) from an audio file.
-    Return sample_index of start of each chunk and original sr
+    audio_file: str, chunk_size: int, sample_rate: int
+) -> List[torch.Tensor]:
+    """Create sequential chunks of size chunk_size from an audio file.
+    Return each chunk
     """
     chunks = []
     audio, sr = torchaudio.load(audio_file)
@@ -138,8 +138,31 @@ def create_sequential_chunks(
     for start in chunk_starts:
         if start + chunk_size > audio.shape[-1]:
             break
-        chunks.append(audio[:, start : start + chunk_size])
-    return chunks, sr
+        chunk = audio[:, start : start + chunk_size]
+        resampled_chunk = torchaudio.functional.resample(chunk, sr, sample_rate)
+        # Skip chunks that are too short
+        if resampled_chunk.shape[-1] < chunk_size:
+            continue
+        chunks.append(chunk)
+    return chunks
+
+
+def select_random_chunk(
+    audio_file: str, chunk_size: int, sample_rate: int
+) -> List[torch.Tensor]:
+    """Select random chunk of size chunk_size (samples) from an audio file."""
+    audio, sr = torchaudio.load(audio_file)
+    new_chunk_size = int(chunk_size * (sr / sample_rate))
+    if new_chunk_size >= audio.shape[-1]:
+        return None
+    max_len = audio.shape[-1] - new_chunk_size
+    random_start = torch.randint(0, max_len, (1,)).item()
+    chunk = audio[:, random_start : random_start + new_chunk_size]
+    # Skip if energy too low
+    if torch.mean(torch.abs(chunk)) < 1e-6:
+        return None
+    resampled_chunk = torchaudio.functional.resample(chunk, sr, sample_rate)
+    return resampled_chunk
 
 
 def spectrogram(

scripts/download.py

@@ -1,8 +1,6 @@
 import os
-import sys
-import glob
-import torch
 import argparse
+import shutil
 
 
 def download_zip_dataset(dataset_url: str, output_dir: str):
@@ -26,8 +24,42 @@ def process_dataset(dataset_dir: str, output_dir: str):
         pass
     elif dataset_dir == "IDMT-SMT-DRUMS-V2":
         pass
+    elif dataset_dir == "DSD100":
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Mixtures"))
+        for dir in os.listdir(os.path.join(output_dir, dataset_dir, "Sources", "Dev")):
+            source = os.path.join(output_dir, dataset_dir, "Sources", "Dev", dir)
+            shutil.move(source, os.path.join(output_dir, dataset_dir))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources", "Dev"))
+        for dir in os.listdir(os.path.join(output_dir, dataset_dir, "Sources", "Test")):
+            source = os.path.join(output_dir, dataset_dir, "Sources", "Test", dir)
+            shutil.move(source, os.path.join(output_dir, dataset_dir))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources", "Test"))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources"))
+
+        os.mkdir(os.path.join(output_dir, dataset_dir, "train"))
+        os.mkdir(os.path.join(output_dir, dataset_dir, "val"))
+        os.mkdir(os.path.join(output_dir, dataset_dir, "test"))
+        files = os.listdir(os.path.join(output_dir, dataset_dir))
+
+        num = 0
+        for dir in files:
+            if not os.path.isdir(os.path.join(output_dir, dataset_dir, dir)):
+                continue
+            if dir == "train" or dir == "val" or dir == "test":
+                continue
+            source = os.path.join(output_dir, dataset_dir, dir, "bass.wav")
+            if num < 80:
+                dest = os.path.join(output_dir, dataset_dir, "train", f"{num}.wav")
+            elif num < 90:
+                dest = os.path.join(output_dir, dataset_dir, "val", f"{num}.wav")
+            else:
+                dest = os.path.join(output_dir, dataset_dir, "test", f"{num}.wav")
+            shutil.move(source, dest)
+            shutil.rmtree(os.path.join(output_dir, dataset_dir, dir))
+            num += 1
+
     else:
-        raise NotImplemented(f"Invalid dataset_dir = {dataset_dir}.")
+        raise NotImplementedError(f"Invalid dataset_dir = {dataset_dir}.")
 
 
 if __name__ == "__main__":
@@ -38,7 +70,7 @@ if __name__ == "__main__":
             "vocalset",
             "guitarset",
             "idmt-smt-guitar",
-            "idmt-smt-bass",
+            "dsd100",
             "idmt-smt-drums",
         ],
         nargs="+",
@@ -49,10 +81,11 @@ if __name__ == "__main__":
         "vocalset": "https://zenodo.org/record/1442513/files/VocalSet1-2.zip",
         "guitarset": "https://zenodo.org/record/3371780/files/audio_mono-mic.zip",
         "IDMT-SMT-GUITAR_V2": "https://zenodo.org/record/7544110/files/IDMT-SMT-GUITAR_V2.zip",
-        "IDMT-SMT-BASS": "https://zenodo.org/record/7188892/files/IDMT-SMT-BASS.zip",
+        "DSD100": "http://liutkus.net/DSD100.zip",
         "IDMT-SMT-DRUMS-V2": "https://zenodo.org/record/7544164/files/IDMT-SMT-DRUMS-V2.zip",
     }
 
     for dataset_name, dataset_url in dataset_urls.items():
         if dataset_name in args.dataset_names:
             download_zip_dataset(dataset_url, "~/data/remfx-data")
+            process_dataset(dataset_name, "~/data/remfx-data")

setup.py

@@ -48,6 +48,7 @@ setup(
         "pedalboard",
         "frechet_audio_distance",
         "ordered-set",
+        "asteroid",
     ],
     include_package_data=True,
     license="Apache License 2.0",

shell_vars.sh

@@ -1,3 +1,3 @@
-export DATASET_ROOT="./data/VocalSet"
+export DATASET_ROOT="./data/remfx-data"
 export WANDB_PROJECT="RemFX"
 export WANDB_ENTITY="mattricesound"