init

.gitattributes

@@ -0,0 +1,29 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+example.wav filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,13 @@
+Copyright 2021 ByteDance
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

README.md

@@ -0,0 +1,33 @@
+---
+title: Music_Source_Separation
+emoji: ⚡
+colorFrom: green
+colorTo: yellow
+sdk: gradio
+app_file: app.py
+pinned: false
+---
+
+# Configuration
+
+`title`: _string_  
+Display title for the Space
+
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,40 @@
+import os
+os.system('pip install gradio==2.3.0a0')
+os.system('pip freeze')
+import sys
+sys.path.append('.')
+import gradio as gr
+os.system('pip install -U torchtext==0.8.0')
+#os.system('python setup.py install --install-dir .')
+from scipy.io import wavfile
+
+os.system('./separate_scripts/download_checkpoints.sh')
+
+def inference(audio):
+    # read the file and get the sample rate and data
+    rate, data = wavfile.read(audio.name) 
+    
+    # save the result
+    wavfile.write('foo_left.wav', rate, data)
+    os.system("""python bytesep/inference.py --config_yaml=./scripts/4_train/musdb18/configs/vocals-accompaniment,resunet_subbandtime.yaml --checkpoint_path=./downloaded_checkpoints/resunet143_subbtandtime_vocals_8.8dB_350k_steps.pth --audio_path=foo_left.wav --output_path=sep_vocals.mp3""")
+    #os.system('./separate_scripts/separate_vocals.sh ' + audio.name + ' "sep_vocals.mp3"')
+    os.system("""python bytesep/inference.py --config_yaml=./scripts/4_train/musdb18/configs/accompaniment-vocals,resunet_subbandtime.yaml --checkpoint_path=./downloaded_checkpoints/resunet143_subbtandtime_accompaniment_16.4dB_350k_steps.pth --audio_path=foo_left.wav --output_path=sep_accompaniment.mp3""")
+    #os.system('./separate_scripts/separate_accompaniment.sh ' + audio.name + ' "sep_accompaniment.mp3"')
+    #os.system('python separate_scripts/separate.py --audio_path=' +audio.name+' --source_type="accompaniment"')
+    #os.system('python separate_scripts/separate.py --audio_path=' +audio.name+' --source_type="vocals"')
+    return 'sep_vocals.mp3', 'sep_accompaniment.mp3'
+title = "Music Source Separation"
+description = "Gradio demo for Music Source Separation. To use it, simply add your audio, or click one of the examples to load them. Currently supports .wav files. Read more at the links below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.05418'>Decoupling Magnitude and Phase Estimation with Deep ResUNet for Music Source Separation</a> | <a href='https://github.com/bytedance/music_source_separation'>Github Repo</a></p>"
+
+examples = [['example.wav']]
+gr.Interface(
+    inference, 
+    gr.inputs.Audio(type="file", label="Input"), 
+    [gr.outputs.Audio(type="file", label="Vocals"),gr.outputs.Audio(type="file", label="Accompaniment")],
+    title=title,
+    description=description,
+    article=article,
+    enable_queue=True,
+    examples=examples
+    ).launch(debug=True)

bytesep/__init__.py

@@ -0,0 +1 @@
+from bytesep.inference import Separator

bytesep/callbacks/__init__.py

@@ -0,0 +1,76 @@
+from typing import List
+
+import pytorch_lightning as pl
+import torch.nn as nn
+
+
+def get_callbacks(
+    task_name: str,
+    config_yaml: str,
+    workspace: str,
+    checkpoints_dir: str,
+    statistics_path: str,
+    logger: pl.loggers.TensorBoardLogger,
+    model: nn.Module,
+    evaluate_device: str,
+) -> List[pl.Callback]:
+    r"""Get callbacks of a task and config yaml file.
+
+    Args:
+        task_name: str
+        config_yaml: str
+        dataset_dir: str
+        workspace: str, containing useful files such as audios for evaluation
+        checkpoints_dir: str, directory to save checkpoints
+        statistics_dir: str, directory to save statistics
+        logger: pl.loggers.TensorBoardLogger
+        model: nn.Module
+        evaluate_device: str
+
+    Return:
+        callbacks: List[pl.Callback]
+    """
+    if task_name == 'musdb18':
+
+        from bytesep.callbacks.musdb18 import get_musdb18_callbacks
+
+        return get_musdb18_callbacks(
+            config_yaml=config_yaml,
+            workspace=workspace,
+            checkpoints_dir=checkpoints_dir,
+            statistics_path=statistics_path,
+            logger=logger,
+            model=model,
+            evaluate_device=evaluate_device,
+        )
+
+    elif task_name == 'voicebank-demand':
+
+        from bytesep.callbacks.voicebank_demand import get_voicebank_demand_callbacks
+
+        return get_voicebank_demand_callbacks(
+            config_yaml=config_yaml,
+            workspace=workspace,
+            checkpoints_dir=checkpoints_dir,
+            statistics_path=statistics_path,
+            logger=logger,
+            model=model,
+            evaluate_device=evaluate_device,
+        )
+
+    elif task_name in ['vctk-musdb18', 'violin-piano', 'piano-symphony']:
+
+        from bytesep.callbacks.instruments_callbacks import get_instruments_callbacks
+
+        return get_instruments_callbacks(
+            config_yaml=config_yaml,
+            workspace=workspace,
+            checkpoints_dir=checkpoints_dir,
+            statistics_path=statistics_path,
+            logger=logger,
+            model=model,
+            evaluate_device=evaluate_device,
+        )
+
+    else:
+        raise NotImplementedError

bytesep/callbacks/base_callbacks.py

@@ -0,0 +1,44 @@
+import logging
+import os
+from typing import NoReturn
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from pytorch_lightning.utilities import rank_zero_only
+
+
+class SaveCheckpointsCallback(pl.Callback):
+    def __init__(
+        self,
+        model: nn.Module,
+        checkpoints_dir: str,
+        save_step_frequency: int,
+    ):
+        r"""Callback to save checkpoints every #save_step_frequency steps.
+
+        Args:
+            model: nn.Module
+            checkpoints_dir: str, directory to save checkpoints
+            save_step_frequency: int
+        """
+        self.model = model
+        self.checkpoints_dir = checkpoints_dir
+        self.save_step_frequency = save_step_frequency
+        os.makedirs(self.checkpoints_dir, exist_ok=True)
+
+    @rank_zero_only
+    def on_batch_end(self, trainer: pl.Trainer, _) -> NoReturn:
+        r"""Save checkpoint."""
+        global_step = trainer.global_step
+
+        if global_step % self.save_step_frequency == 0:
+
+            checkpoint_path = os.path.join(
+                self.checkpoints_dir, "step={}.pth".format(global_step)
+            )
+
+            checkpoint = {'step': global_step, 'model': self.model.state_dict()}
+
+            torch.save(checkpoint, checkpoint_path)
+            logging.info("Save checkpoint to {}".format(checkpoint_path))

bytesep/callbacks/instruments_callbacks.py

@@ -0,0 +1,200 @@
+import logging
+import os
+import time
+from typing import List, NoReturn
+
+import librosa
+import numpy as np
+import pytorch_lightning as pl
+import torch.nn as nn
+from pytorch_lightning.utilities import rank_zero_only
+
+from bytesep.callbacks.base_callbacks import SaveCheckpointsCallback
+from bytesep.inference import Separator
+from bytesep.utils import StatisticsContainer, calculate_sdr, read_yaml
+
+
+def get_instruments_callbacks(
+    config_yaml: str,
+    workspace: str,
+    checkpoints_dir: str,
+    statistics_path: str,
+    logger: pl.loggers.TensorBoardLogger,
+    model: nn.Module,
+    evaluate_device: str,
+) -> List[pl.Callback]:
+    """Get Voicebank-Demand callbacks of a config yaml.
+
+    Args:
+        config_yaml: str
+        workspace: str
+        checkpoints_dir: str, directory to save checkpoints
+        statistics_dir: str, directory to save statistics
+        logger: pl.loggers.TensorBoardLogger
+        model: nn.Module
+        evaluate_device: str
+
+    Return:
+        callbacks: List[pl.Callback]
+    """
+    configs = read_yaml(config_yaml)
+    task_name = configs['task_name']
+    target_source_types = configs['train']['target_source_types']
+    input_channels = configs['train']['channels']
+    mono = True if input_channels == 1 else False
+    test_audios_dir = os.path.join(workspace, "evaluation_audios", task_name, "test")
+    sample_rate = configs['train']['sample_rate']
+    evaluate_step_frequency = configs['train']['evaluate_step_frequency']
+    save_step_frequency = configs['train']['save_step_frequency']
+    test_batch_size = configs['evaluate']['batch_size']
+    test_segment_seconds = configs['evaluate']['segment_seconds']
+
+    test_segment_samples = int(test_segment_seconds * sample_rate)
+    assert len(target_source_types) == 1
+    target_source_type = target_source_types[0]
+
+    # save checkpoint callback
+    save_checkpoints_callback = SaveCheckpointsCallback(
+        model=model,
+        checkpoints_dir=checkpoints_dir,
+        save_step_frequency=save_step_frequency,
+    )
+
+    # statistics container
+    statistics_container = StatisticsContainer(statistics_path)
+
+    # evaluation callback
+    evaluate_test_callback = EvaluationCallback(
+        model=model,
+        target_source_type=target_source_type,
+        input_channels=input_channels,
+        sample_rate=sample_rate,
+        mono=mono,
+        evaluation_audios_dir=test_audios_dir,
+        segment_samples=test_segment_samples,
+        batch_size=test_batch_size,
+        device=evaluate_device,
+        evaluate_step_frequency=evaluate_step_frequency,
+        logger=logger,
+        statistics_container=statistics_container,
+    )
+
+    callbacks = [save_checkpoints_callback, evaluate_test_callback]
+    # callbacks = [save_checkpoints_callback]
+
+    return callbacks
+
+
+class EvaluationCallback(pl.Callback):
+    def __init__(
+        self,
+        model: nn.Module,
+        input_channels: int,
+        evaluation_audios_dir: str,
+        target_source_type: str,
+        sample_rate: int,
+        mono: bool,
+        segment_samples: int,
+        batch_size: int,
+        device: str,
+        evaluate_step_frequency: int,
+        logger: pl.loggers.TensorBoardLogger,
+        statistics_container: StatisticsContainer,
+    ):
+        r"""Callback to evaluate every #save_step_frequency steps.
+
+        Args:
+            model: nn.Module
+            input_channels: int
+            evaluation_audios_dir: str, directory containing audios for evaluation
+            target_source_type: str, e.g., 'violin'
+            sample_rate: int
+            mono: bool
+            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
+            batch_size, int, e.g., 12
+            device: str, e.g., 'cuda'
+            evaluate_step_frequency: int, evaluate every #save_step_frequency steps
+            logger: pl.loggers.TensorBoardLogger
+            statistics_container: StatisticsContainer
+        """
+        self.model = model
+        self.target_source_type = target_source_type
+        self.sample_rate = sample_rate
+        self.mono = mono
+        self.segment_samples = segment_samples
+        self.evaluate_step_frequency = evaluate_step_frequency
+        self.logger = logger
+        self.statistics_container = statistics_container
+
+        self.evaluation_audios_dir = evaluation_audios_dir
+
+        # separator
+        self.separator = Separator(model, self.segment_samples, batch_size, device)
+
+    @rank_zero_only
+    def on_batch_end(self, trainer: pl.Trainer, _) -> NoReturn:
+        r"""Evaluate losses on a few mini-batches. Losses are only used for
+        observing training, and are not final F1 metrics.
+        """
+
+        global_step = trainer.global_step
+
+        if global_step % self.evaluate_step_frequency == 0:
+
+            mixture_audios_dir = os.path.join(self.evaluation_audios_dir, 'mixture')
+            clean_audios_dir = os.path.join(
+                self.evaluation_audios_dir, self.target_source_type
+            )
+
+            audio_names = sorted(os.listdir(mixture_audios_dir))
+
+            error_str = "Directory {} does not contain audios for evaluation!".format(
+                self.evaluation_audios_dir
+            )
+            assert len(audio_names) > 0, error_str
+
+            logging.info("--- Step {} ---".format(global_step))
+            logging.info("Total {} pieces for evaluation:".format(len(audio_names)))
+
+            eval_time = time.time()
+
+            sdrs = []
+
+            for n, audio_name in enumerate(audio_names):
+
+                # Load audio.
+                mixture_path = os.path.join(mixture_audios_dir, audio_name)
+                clean_path = os.path.join(clean_audios_dir, audio_name)
+
+                mixture, origin_fs = librosa.core.load(
+                    mixture_path, sr=self.sample_rate, mono=self.mono
+                )
+
+                # Target
+                clean, origin_fs = librosa.core.load(
+                    clean_path, sr=self.sample_rate, mono=self.mono
+                )
+
+                if mixture.ndim == 1:
+                    mixture = mixture[None, :]
+                # (channels_num, audio_length)
+
+                input_dict = {'waveform': mixture}
+
+                # separate
+                sep_wav = self.separator.separate(input_dict)
+                # (channels_num, audio_length)
+
+                sdr = calculate_sdr(ref=clean, est=sep_wav)
+
+                print("{} SDR: {:.3f}".format(audio_name, sdr))
+                sdrs.append(sdr)
+
+            logging.info("-----------------------------")
+            logging.info('Avg SDR: {:.3f}'.format(np.mean(sdrs)))
+
+            logging.info("Evlauation time: {:.3f}".format(time.time() - eval_time))
+
+            statistics = {"sdr": np.mean(sdrs)}
+            self.statistics_container.append(global_step, statistics, 'test')
+            self.statistics_container.dump()

bytesep/callbacks/musdb18.py

@@ -0,0 +1,485 @@
+import logging
+import os
+import time
+from typing import Dict, List, NoReturn
+
+import librosa
+import musdb
+import museval
+import numpy as np
+import pytorch_lightning as pl
+import torch.nn as nn
+from pytorch_lightning.utilities import rank_zero_only
+
+from bytesep.callbacks.base_callbacks import SaveCheckpointsCallback
+from bytesep.dataset_creation.pack_audios_to_hdf5s.musdb18 import preprocess_audio
+from bytesep.inference import Separator
+from bytesep.utils import StatisticsContainer, read_yaml
+
+
+def get_musdb18_callbacks(
+    config_yaml: str,
+    workspace: str,
+    checkpoints_dir: str,
+    statistics_path: str,
+    logger: pl.loggers.TensorBoardLogger,
+    model: nn.Module,
+    evaluate_device: str,
+) -> List[pl.Callback]:
+    r"""Get MUSDB18 callbacks of a config yaml.
+
+    Args:
+        config_yaml: str
+        workspace: str
+        checkpoints_dir: str, directory to save checkpoints
+        statistics_dir: str, directory to save statistics
+        logger: pl.loggers.TensorBoardLogger
+        model: nn.Module
+        evaluate_device: str
+
+    Return:
+        callbacks: List[pl.Callback]
+    """
+    configs = read_yaml(config_yaml)
+    task_name = configs['task_name']
+    evaluation_callback = configs['train']['evaluation_callback']
+    target_source_types = configs['train']['target_source_types']
+    input_channels = configs['train']['channels']
+    evaluation_audios_dir = os.path.join(workspace, "evaluation_audios", task_name)
+    test_segment_seconds = configs['evaluate']['segment_seconds']
+    sample_rate = configs['train']['sample_rate']
+    test_segment_samples = int(test_segment_seconds * sample_rate)
+    test_batch_size = configs['evaluate']['batch_size']
+
+    evaluate_step_frequency = configs['train']['evaluate_step_frequency']
+    save_step_frequency = configs['train']['save_step_frequency']
+
+    # save checkpoint callback
+    save_checkpoints_callback = SaveCheckpointsCallback(
+        model=model,
+        checkpoints_dir=checkpoints_dir,
+        save_step_frequency=save_step_frequency,
+    )
+
+    # evaluation callback
+    EvaluationCallback = _get_evaluation_callback_class(evaluation_callback)
+
+    # statistics container
+    statistics_container = StatisticsContainer(statistics_path)
+
+    # evaluation callback
+    evaluate_train_callback = EvaluationCallback(
+        dataset_dir=evaluation_audios_dir,
+        model=model,
+        target_source_types=target_source_types,
+        input_channels=input_channels,
+        sample_rate=sample_rate,
+        split='train',
+        segment_samples=test_segment_samples,
+        batch_size=test_batch_size,
+        device=evaluate_device,
+        evaluate_step_frequency=evaluate_step_frequency,
+        logger=logger,
+        statistics_container=statistics_container,
+    )
+
+    evaluate_test_callback = EvaluationCallback(
+        dataset_dir=evaluation_audios_dir,
+        model=model,
+        target_source_types=target_source_types,
+        input_channels=input_channels,
+        sample_rate=sample_rate,
+        split='test',
+        segment_samples=test_segment_samples,
+        batch_size=test_batch_size,
+        device=evaluate_device,
+        evaluate_step_frequency=evaluate_step_frequency,
+        logger=logger,
+        statistics_container=statistics_container,
+    )
+
+    # callbacks = [save_checkpoints_callback, evaluate_train_callback, evaluate_test_callback]
+    callbacks = [save_checkpoints_callback, evaluate_test_callback]
+
+    return callbacks
+
+
+def _get_evaluation_callback_class(evaluation_callback) -> pl.Callback:
+    r"""Get evaluation callback class."""
+    if evaluation_callback == "Musdb18EvaluationCallback":
+        return Musdb18EvaluationCallback
+
+    if evaluation_callback == 'Musdb18ConditionalEvaluationCallback':
+        return Musdb18ConditionalEvaluationCallback
+
+    else:
+        raise NotImplementedError
+
+
+class Musdb18EvaluationCallback(pl.Callback):
+    def __init__(
+        self,
+        dataset_dir: str,
+        model: nn.Module,
+        target_source_types: str,
+        input_channels: int,
+        split: str,
+        sample_rate: int,
+        segment_samples: int,
+        batch_size: int,
+        device: str,
+        evaluate_step_frequency: int,
+        logger: pl.loggers.TensorBoardLogger,
+        statistics_container: StatisticsContainer,
+    ):
+        r"""Callback to evaluate every #save_step_frequency steps.
+
+        Args:
+            dataset_dir: str
+            model: nn.Module
+            target_source_types: List[str], e.g., ['vocals', 'bass', ...]
+            input_channels: int
+            split: 'train' | 'test'
+            sample_rate: int
+            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
+            batch_size, int, e.g., 12
+            device: str, e.g., 'cuda'
+            evaluate_step_frequency: int, evaluate every #save_step_frequency steps
+            logger: object
+            statistics_container: StatisticsContainer
+        """
+        self.model = model
+        self.target_source_types = target_source_types
+        self.input_channels = input_channels
+        self.sample_rate = sample_rate
+        self.split = split
+        self.segment_samples = segment_samples
+        self.evaluate_step_frequency = evaluate_step_frequency
+        self.logger = logger
+        self.statistics_container = statistics_container
+        self.mono = input_channels == 1
+        self.resample_type = "kaiser_fast"
+
+        self.mus = musdb.DB(root=dataset_dir, subsets=[split])
+
+        error_msg = "The directory {} is empty!".format(dataset_dir)
+        assert len(self.mus) > 0, error_msg
+
+        # separator
+        self.separator = Separator(model, self.segment_samples, batch_size, device)
+
+    @rank_zero_only
+    def on_batch_end(self, trainer: pl.Trainer, _) -> NoReturn:
+        r"""Evaluate separation SDRs of audio recordings."""
+        global_step = trainer.global_step
+
+        if global_step % self.evaluate_step_frequency == 0:
+
+            sdr_dict = {}
+
+            logging.info("--- Step {} ---".format(global_step))
+            logging.info("Total {} pieces for evaluation:".format(len(self.mus.tracks)))
+
+            eval_time = time.time()
+
+            for track in self.mus.tracks:
+
+                audio_name = track.name
+
+                # Get waveform of mixture.
+                mixture = track.audio.T
+                # (channels_num, audio_samples)
+
+                mixture = preprocess_audio(
+                    audio=mixture,
+                    mono=self.mono,
+                    origin_sr=track.rate,
+                    sr=self.sample_rate,
+                    resample_type=self.resample_type,
+                )
+                # (channels_num, audio_samples)
+
+                target_dict = {}
+                sdr_dict[audio_name] = {}
+
+                # Get waveform of all target source types.
+                for j, source_type in enumerate(self.target_source_types):
+                    # E.g., ['vocals', 'bass', ...]
+
+                    audio = track.targets[source_type].audio.T
+
+                    audio = preprocess_audio(
+                        audio=audio,
+                        mono=self.mono,
+                        origin_sr=track.rate,
+                        sr=self.sample_rate,
+                        resample_type=self.resample_type,
+                    )
+                    # (channels_num, audio_samples)
+
+                    target_dict[source_type] = audio
+                    # (channels_num, audio_samples)
+
+                # Separate.
+                input_dict = {'waveform': mixture}
+
+                sep_wavs = self.separator.separate(input_dict)
+                # sep_wavs: (target_sources_num * channels_num, audio_samples)
+
+                # Post process separation results.
+                sep_wavs = preprocess_audio(
+                    audio=sep_wavs,
+                    mono=self.mono,
+                    origin_sr=self.sample_rate,
+                    sr=track.rate,
+                    resample_type=self.resample_type,
+                )
+                # sep_wavs: (target_sources_num * channels_num, audio_samples)
+
+                sep_wavs = librosa.util.fix_length(
+                    sep_wavs, size=mixture.shape[1], axis=1
+                )
+                # sep_wavs: (target_sources_num * channels_num, audio_samples)
+
+                sep_wav_dict = get_separated_wavs_from_simo_output(
+                    sep_wavs, self.input_channels, self.target_source_types
+                )
+                # output_dict: dict, e.g., {
+                #     'vocals': (channels_num, audio_samples),
+                #     'bass': (channels_num, audio_samples),
+                #     ...,
+                # }
+
+                # Evaluate for all target source types.
+                for source_type in self.target_source_types:
+                    # E.g., ['vocals', 'bass', ...]
+
+                    # Calculate SDR using museval, input shape should be: (nsrc, nsampl, nchan).
+                    (sdrs, _, _, _) = museval.evaluate(
+                        [target_dict[source_type].T], [sep_wav_dict[source_type].T]
+                    )
+
+                    sdr = np.nanmedian(sdrs)
+                    sdr_dict[audio_name][source_type] = sdr
+
+                    logging.info(
+                        "{}, {}, sdr: {:.3f}".format(audio_name, source_type, sdr)
+                    )
+
+            logging.info("-----------------------------")
+            median_sdr_dict = {}
+
+            # Calculate median SDRs of all songs.
+            for source_type in self.target_source_types:
+                # E.g., ['vocals', 'bass', ...]
+
+                median_sdr = np.median(
+                    [
+                        sdr_dict[audio_name][source_type]
+                        for audio_name in sdr_dict.keys()
+                    ]
+                )
+
+                median_sdr_dict[source_type] = median_sdr
+
+                logging.info(
+                    "Step: {}, {}, Median SDR: {:.3f}".format(
+                        global_step, source_type, median_sdr
+                    )
+                )
+
+            logging.info("Evlauation time: {:.3f}".format(time.time() - eval_time))
+
+            statistics = {"sdr_dict": sdr_dict, "median_sdr_dict": median_sdr_dict}
+            self.statistics_container.append(global_step, statistics, self.split)
+            self.statistics_container.dump()
+
+
+def get_separated_wavs_from_simo_output(x, input_channels, target_source_types) -> Dict:
+    r"""Get separated waveforms of target sources from a single input multiple
+    output (SIMO) system.
+
+    Args:
+        x: (target_sources_num * channels_num, audio_samples)
+        input_channels: int
+        target_source_types: List[str], e.g., ['vocals', 'bass', ...]
+
+    Returns:
+        output_dict: dict, e.g., {
+            'vocals': (channels_num, audio_samples),
+            'bass': (channels_num, audio_samples),
+            ...,
+        }
+    """
+    output_dict = {}
+
+    for j, source_type in enumerate(target_source_types):
+        output_dict[source_type] = x[j * input_channels : (j + 1) * input_channels]
+
+    return output_dict
+
+
+class Musdb18ConditionalEvaluationCallback(pl.Callback):
+    def __init__(
+        self,
+        dataset_dir: str,
+        model: nn.Module,
+        target_source_types: str,
+        input_channels: int,
+        split: str,
+        sample_rate: int,
+        segment_samples: int,
+        batch_size: int,
+        device: str,
+        evaluate_step_frequency: int,
+        logger: pl.loggers.TensorBoardLogger,
+        statistics_container: StatisticsContainer,
+    ):
+        r"""Callback to evaluate every #save_step_frequency steps.
+
+        Args:
+            dataset_dir: str
+            model: nn.Module
+            target_source_types: List[str], e.g., ['vocals', 'bass', ...]
+            input_channels: int
+            split: 'train' | 'test'
+            sample_rate: int
+            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
+            batch_size, int, e.g., 12
+            device: str, e.g., 'cuda'
+            evaluate_step_frequency: int, evaluate every #save_step_frequency steps
+            logger: object
+            statistics_container: StatisticsContainer
+        """
+        self.model = model
+        self.target_source_types = target_source_types
+        self.input_channels = input_channels
+        self.sample_rate = sample_rate
+        self.split = split
+        self.segment_samples = segment_samples
+        self.evaluate_step_frequency = evaluate_step_frequency
+        self.logger = logger
+        self.statistics_container = statistics_container
+        self.mono = input_channels == 1
+        self.resample_type = "kaiser_fast"
+
+        self.mus = musdb.DB(root=dataset_dir, subsets=[split])
+
+        error_msg = "The directory {} is empty!".format(dataset_dir)
+        assert len(self.mus) > 0, error_msg
+
+        # separator
+        self.separator = Separator(model, self.segment_samples, batch_size, device)
+
+    @rank_zero_only
+    def on_batch_end(self, trainer: pl.Trainer, _) -> NoReturn:
+        r"""Evaluate separation SDRs of audio recordings."""
+        global_step = trainer.global_step
+
+        if global_step % self.evaluate_step_frequency == 0:
+
+            sdr_dict = {}
+
+            logging.info("--- Step {} ---".format(global_step))
+            logging.info("Total {} pieces for evaluation:".format(len(self.mus.tracks)))
+
+            eval_time = time.time()
+
+            for track in self.mus.tracks:
+
+                audio_name = track.name
+
+                # Get waveform of mixture.
+                mixture = track.audio.T
+                # (channels_num, audio_samples)
+
+                mixture = preprocess_audio(
+                    audio=mixture,
+                    mono=self.mono,
+                    origin_sr=track.rate,
+                    sr=self.sample_rate,
+                    resample_type=self.resample_type,
+                )
+                # (channels_num, audio_samples)
+
+                target_dict = {}
+                sdr_dict[audio_name] = {}
+
+                # Get waveform of all target source types.
+                for j, source_type in enumerate(self.target_source_types):
+                    # E.g., ['vocals', 'bass', ...]
+
+                    audio = track.targets[source_type].audio.T
+
+                    audio = preprocess_audio(
+                        audio=audio,
+                        mono=self.mono,
+                        origin_sr=track.rate,
+                        sr=self.sample_rate,
+                        resample_type=self.resample_type,
+                    )
+                    # (channels_num, audio_samples)
+
+                    target_dict[source_type] = audio
+                    # (channels_num, audio_samples)
+
+                    condition = np.zeros(len(self.target_source_types))
+                    condition[j] = 1
+
+                    input_dict = {'waveform': mixture, 'condition': condition}
+
+                    sep_wav = self.separator.separate(input_dict)
+                    # sep_wav: (channels_num, audio_samples)
+
+                    sep_wav = preprocess_audio(
+                        audio=sep_wav,
+                        mono=self.mono,
+                        origin_sr=self.sample_rate,
+                        sr=track.rate,
+                        resample_type=self.resample_type,
+                    )
+                    # sep_wav: (channels_num, audio_samples)
+
+                    sep_wav = librosa.util.fix_length(
+                        sep_wav, size=mixture.shape[1], axis=1
+                    )
+                    # sep_wav: (target_sources_num * channels_num, audio_samples)
+
+                    # Calculate SDR using museval, input shape should be: (nsrc, nsampl, nchan)
+                    (sdrs, _, _, _) = museval.evaluate(
+                        [target_dict[source_type].T], [sep_wav.T]
+                    )
+
+                    sdr = np.nanmedian(sdrs)
+                    sdr_dict[audio_name][source_type] = sdr
+
+                    logging.info(
+                        "{}, {}, sdr: {:.3f}".format(audio_name, source_type, sdr)
+                    )
+
+            logging.info("-----------------------------")
+            median_sdr_dict = {}
+
+            # Calculate median SDRs of all songs.
+            for source_type in self.target_source_types:
+
+                median_sdr = np.median(
+                    [
+                        sdr_dict[audio_name][source_type]
+                        for audio_name in sdr_dict.keys()
+                    ]
+                )
+
+                median_sdr_dict[source_type] = median_sdr
+
+                logging.info(
+                    "Step: {}, {}, Median SDR: {:.3f}".format(
+                        global_step, source_type, median_sdr
+                    )
+                )
+
+            logging.info("Evlauation time: {:.3f}".format(time.time() - eval_time))
+
+            statistics = {"sdr_dict": sdr_dict, "median_sdr_dict": median_sdr_dict}
+            self.statistics_container.append(global_step, statistics, self.split)
+            self.statistics_container.dump()

bytesep/callbacks/voicebank_demand.py

@@ -0,0 +1,231 @@
+import logging
+import os
+import time
+from typing import List, NoReturn
+
+import librosa
+import numpy as np
+import pysepm
+import pytorch_lightning as pl
+import torch.nn as nn
+from pesq import pesq
+from pytorch_lightning.utilities import rank_zero_only
+
+from bytesep.callbacks.base_callbacks import SaveCheckpointsCallback
+from bytesep.inference import Separator
+from bytesep.utils import StatisticsContainer, read_yaml
+
+
+def get_voicebank_demand_callbacks(
+    config_yaml: str,
+    workspace: str,
+    checkpoints_dir: str,
+    statistics_path: str,
+    logger: pl.loggers.TensorBoardLogger,
+    model: nn.Module,
+    evaluate_device: str,
+) -> List[pl.Callback]:
+    """Get Voicebank-Demand callbacks of a config yaml.
+
+    Args:
+        config_yaml: str
+        workspace: str
+        checkpoints_dir: str, directory to save checkpoints
+        statistics_dir: str, directory to save statistics
+        logger: pl.loggers.TensorBoardLogger
+        model: nn.Module
+        evaluate_device: str
+
+    Return:
+        callbacks: List[pl.Callback]
+    """
+    configs = read_yaml(config_yaml)
+    task_name = configs['task_name']
+    target_source_types = configs['train']['target_source_types']
+    input_channels = configs['train']['channels']
+    evaluation_audios_dir = os.path.join(workspace, "evaluation_audios", task_name)
+    sample_rate = configs['train']['sample_rate']
+    evaluate_step_frequency = configs['train']['evaluate_step_frequency']
+    save_step_frequency = configs['train']['save_step_frequency']
+    test_batch_size = configs['evaluate']['batch_size']
+    test_segment_seconds = configs['evaluate']['segment_seconds']
+
+    test_segment_samples = int(test_segment_seconds * sample_rate)
+    assert len(target_source_types) == 1
+    target_source_type = target_source_types[0]
+    assert target_source_type == 'speech'
+
+    # save checkpoint callback
+    save_checkpoints_callback = SaveCheckpointsCallback(
+        model=model,
+        checkpoints_dir=checkpoints_dir,
+        save_step_frequency=save_step_frequency,
+    )
+
+    # statistics container
+    statistics_container = StatisticsContainer(statistics_path)
+
+    # evaluation callback
+    evaluate_test_callback = EvaluationCallback(
+        model=model,
+        input_channels=input_channels,
+        sample_rate=sample_rate,
+        evaluation_audios_dir=evaluation_audios_dir,
+        segment_samples=test_segment_samples,
+        batch_size=test_batch_size,
+        device=evaluate_device,
+        evaluate_step_frequency=evaluate_step_frequency,
+        logger=logger,
+        statistics_container=statistics_container,
+    )
+
+    callbacks = [save_checkpoints_callback, evaluate_test_callback]
+
+    return callbacks
+
+
+class EvaluationCallback(pl.Callback):
+    def __init__(
+        self,
+        model: nn.Module,
+        input_channels: int,
+        evaluation_audios_dir,
+        sample_rate: int,
+        segment_samples: int,
+        batch_size: int,
+        device: str,
+        evaluate_step_frequency: int,
+        logger: pl.loggers.TensorBoardLogger,
+        statistics_container: StatisticsContainer,
+    ):
+        r"""Callback to evaluate every #save_step_frequency steps.
+
+        Args:
+            model: nn.Module
+            input_channels: int
+            evaluation_audios_dir: str, directory containing audios for evaluation
+            sample_rate: int
+            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
+            batch_size, int, e.g., 12
+            device: str, e.g., 'cuda'
+            evaluate_step_frequency: int, evaluate every #save_step_frequency steps
+            logger: pl.loggers.TensorBoardLogger
+            statistics_container: StatisticsContainer
+        """
+        self.model = model
+        self.mono = True
+        self.sample_rate = sample_rate
+        self.segment_samples = segment_samples
+        self.evaluate_step_frequency = evaluate_step_frequency
+        self.logger = logger
+        self.statistics_container = statistics_container
+
+        self.clean_dir = os.path.join(evaluation_audios_dir, "clean_testset_wav")
+        self.noisy_dir = os.path.join(evaluation_audios_dir, "noisy_testset_wav")
+
+        self.EVALUATION_SAMPLE_RATE = 16000  # Evaluation sample rate of the
+        # Voicebank-Demand task.
+
+        # separator
+        self.separator = Separator(model, self.segment_samples, batch_size, device)
+
+    @rank_zero_only
+    def on_batch_end(self, trainer: pl.Trainer, _) -> NoReturn:
+        r"""Evaluate losses on a few mini-batches. Losses are only used for
+        observing training, and are not final F1 metrics.
+        """
+
+        global_step = trainer.global_step
+
+        if global_step % self.evaluate_step_frequency == 0:
+
+            audio_names = sorted(
+                [
+                    audio_name
+                    for audio_name in sorted(os.listdir(self.clean_dir))
+                    if audio_name.endswith('.wav')
+                ]
+            )
+
+            error_str = "Directory {} does not contain audios for evaluation!".format(
+                self.clean_dir
+            )
+            assert len(audio_names) > 0, error_str
+
+            pesqs, csigs, cbaks, covls, ssnrs = [], [], [], [], []
+
+            logging.info("--- Step {} ---".format(global_step))
+            logging.info("Total {} pieces for evaluation:".format(len(audio_names)))
+
+            eval_time = time.time()
+
+            for n, audio_name in enumerate(audio_names):
+
+                # Load audio.
+                clean_path = os.path.join(self.clean_dir, audio_name)
+                mixture_path = os.path.join(self.noisy_dir, audio_name)
+
+                mixture, _ = librosa.core.load(
+                    mixture_path, sr=self.sample_rate, mono=self.mono
+                )
+
+                if mixture.ndim == 1:
+                    mixture = mixture[None, :]
+                # (channels_num, audio_length)
+
+                # Separate.
+                input_dict = {'waveform': mixture}
+
+                sep_wav = self.separator.separate(input_dict)
+                # (channels_num, audio_length)
+
+                # Target
+                clean, _ = librosa.core.load(
+                    clean_path, sr=self.EVALUATION_SAMPLE_RATE, mono=self.mono
+                )
+
+                # to mono
+                sep_wav = np.squeeze(sep_wav)
+
+                # Resample for evaluation.
+                sep_wav = librosa.resample(
+                    sep_wav,
+                    orig_sr=self.sample_rate,
+                    target_sr=self.EVALUATION_SAMPLE_RATE,
+                )
+
+                sep_wav = librosa.util.fix_length(sep_wav, size=len(clean), axis=0)
+                # (channels, audio_length)
+
+                # Evaluate metrics
+                pesq_ = pesq(self.EVALUATION_SAMPLE_RATE, clean, sep_wav, 'wb')
+
+                (csig, cbak, covl) = pysepm.composite(
+                    clean, sep_wav, self.EVALUATION_SAMPLE_RATE
+                )
+
+                ssnr = pysepm.SNRseg(clean, sep_wav, self.EVALUATION_SAMPLE_RATE)
+
+                pesqs.append(pesq_)
+                csigs.append(csig)
+                cbaks.append(cbak)
+                covls.append(covl)
+                ssnrs.append(ssnr)
+                print(
+                    '{}, {}, PESQ: {:.3f}, CSIG: {:.3f}, CBAK: {:.3f}, COVL: {:.3f}, SSNR: {:.3f}'.format(
+                        n, audio_name, pesq_, csig, cbak, covl, ssnr
+                    )
+                )
+
+            logging.info("-----------------------------")
+            logging.info('Avg PESQ: {:.3f}'.format(np.mean(pesqs)))
+            logging.info('Avg CSIG: {:.3f}'.format(np.mean(csigs)))
+            logging.info('Avg CBAK: {:.3f}'.format(np.mean(cbaks)))
+            logging.info('Avg COVL: {:.3f}'.format(np.mean(covls)))
+            logging.info('Avg SSNR: {:.3f}'.format(np.mean(ssnrs)))
+
+            logging.info("Evlauation time: {:.3f}".format(time.time() - eval_time))
+
+            statistics = {"pesq": np.mean(pesqs)}
+            self.statistics_container.append(global_step, statistics, 'test')
+            self.statistics_container.dump()

bytesep/data/augmentors.py

@@ -0,0 +1,157 @@
+from typing import Dict
+
+import librosa
+import numpy as np
+
+from bytesep.utils import db_to_magnitude, get_pitch_shift_factor, magnitude_to_db
+
+
+class Augmentor:
+    def __init__(self, augmentations: Dict, random_seed=1234):
+        r"""Augmentor for data augmentation of a waveform.
+
+        Args:
+            augmentations: Dict, e.g, {
+                'mixaudio': {'vocals': 2, 'accompaniment': 2}
+                'pitch_shift': {'vocals': 4, 'accompaniment': 4},
+                ...,
+            }
+            random_seed: int
+        """
+        self.augmentations = augmentations
+        self.random_state = np.random.RandomState(random_seed)
+
+    def __call__(self, waveform: np.array, source_type: str) -> np.array:
+        r"""Augment a waveform.
+
+        Args:
+            waveform: (channels_num, audio_samples)
+            source_type: str
+
+        Returns:
+            new_waveform: (channels_num, new_audio_samples)
+        """
+        if 'pitch_shift' in self.augmentations.keys():
+            waveform = self.pitch_shift(waveform, source_type)
+
+        if 'magnitude_scale' in self.augmentations.keys():
+            waveform = self.magnitude_scale(waveform, source_type)
+
+        if 'swap_channel' in self.augmentations.keys():
+            waveform = self.swap_channel(waveform, source_type)
+
+        if 'flip_axis' in self.augmentations.keys():
+            waveform = self.flip_axis(waveform, source_type)
+
+        return waveform
+
+    def pitch_shift(self, waveform: np.array, source_type: str) -> np.array:
+        r"""Shift the pitch of a waveform. We use resampling for fast pitch
+        shifting, so the speed will also be chaneged. The length of the returned
+        waveform will be changed.
+
+        Args:
+            waveform: (channels_num, audio_samples)
+            source_type: str
+
+        Returns:
+            new_waveform: (channels_num, new_audio_samples)
+        """
+
+        # maximum pitch shift in semitones
+        max_pitch_shift = self.augmentations['pitch_shift'][source_type]
+
+        if max_pitch_shift == 0:  # No pitch shift augmentations.
+            return waveform
+
+        # random pitch shift
+        rand_pitch = self.random_state.uniform(
+            low=-max_pitch_shift, high=max_pitch_shift
+        )
+
+        # We use librosa.resample instead of librosa.effects.pitch_shift
+        # because it is 10x times faster.
+        pitch_shift_factor = get_pitch_shift_factor(rand_pitch)
+        dummy_sample_rate = 10000  # Dummy constant.
+
+        channels_num = waveform.shape[0]
+
+        if channels_num == 1:
+            waveform = np.squeeze(waveform)
+
+        new_waveform = librosa.resample(
+            y=waveform,
+            orig_sr=dummy_sample_rate,
+            target_sr=dummy_sample_rate / pitch_shift_factor,
+            res_type='linear',
+            axis=-1,
+        )
+
+        if channels_num == 1:
+            new_waveform = new_waveform[None, :]
+
+        return new_waveform
+
+    def magnitude_scale(self, waveform: np.array, source_type: str) -> np.array:
+        r"""Scale the magnitude of a waveform.
+
+        Args:
+            waveform: (channels_num, audio_samples)
+            source_type: str
+
+        Returns:
+            new_waveform: (channels_num, audio_samples)
+        """
+        lower_db = self.augmentations['magnitude_scale'][source_type]['lower_db']
+        higher_db = self.augmentations['magnitude_scale'][source_type]['higher_db']
+
+        if lower_db == 0 and higher_db == 0:  # No magnitude scale augmentation.
+            return waveform
+
+        # The magnitude (in dB) of the sample with the maximum value.
+        waveform_db = magnitude_to_db(np.max(np.abs(waveform)))
+
+        new_waveform_db = self.random_state.uniform(
+            waveform_db + lower_db, min(waveform_db + higher_db, 0)
+        )
+
+        relative_db = new_waveform_db - waveform_db
+
+        relative_scale = db_to_magnitude(relative_db)
+
+        new_waveform = waveform * relative_scale
+
+        return new_waveform
+
+    def swap_channel(self, waveform: np.array, source_type: str) -> np.array:
+        r"""Randomly swap channels.
+
+        Args:
+            waveform: (channels_num, audio_samples)
+            source_type: str
+
+        Returns:
+            new_waveform: (channels_num, audio_samples)
+        """
+        ndim = waveform.shape[0]
+
+        if ndim == 1:
+            return waveform
+        else:
+            random_axes = self.random_state.permutation(ndim)
+            return waveform[random_axes, :]
+
+    def flip_axis(self, waveform: np.array, source_type: str) -> np.array:
+        r"""Randomly flip the waveform along x-axis.
+
+        Args:
+            waveform: (channels_num, audio_samples)
+            source_type: str
+
+        Returns:
+            new_waveform: (channels_num, audio_samples)
+        """
+        ndim = waveform.shape[0]
+        random_values = self.random_state.choice([-1, 1], size=ndim)
+
+        return waveform * random_values[:, None]

bytesep/data/batch_data_preprocessors.py

@@ -0,0 +1,141 @@
+from typing import Dict, List
+
+import torch
+
+
+class BasicBatchDataPreprocessor:
+    def __init__(self, target_source_types: List[str]):
+        r"""Batch data preprocessor. Used for preparing mixtures and targets for
+        training. If there are multiple target source types, the waveforms of
+        those sources will be stacked along the channel dimension.
+
+        Args:
+            target_source_types: List[str], e.g., ['vocals', 'bass', ...]
+        """
+        self.target_source_types = target_source_types
+
+    def __call__(self, batch_data_dict: Dict) -> List[Dict]:
+        r"""Format waveforms and targets for training.
+
+        Args:
+            batch_data_dict: dict, e.g., {
+                'mixture': (batch_size, channels_num, segment_samples),
+                'vocals': (batch_size, channels_num, segment_samples),
+                'bass': (batch_size, channels_num, segment_samples),
+                ...,
+            }
+
+        Returns:
+            input_dict: dict, e.g., {
+                'waveform': (batch_size, channels_num, segment_samples),
+            }
+            output_dict: dict, e.g., {
+                'target': (batch_size, target_sources_num * channels_num, segment_samples)
+            }
+        """
+        mixtures = batch_data_dict['mixture']
+        # mixtures: (batch_size, channels_num, segment_samples)
+
+        # Concatenate waveforms of multiple targets along the channel axis.
+        targets = torch.cat(
+            [batch_data_dict[source_type] for source_type in self.target_source_types],
+            dim=1,
+        )
+        # targets: (batch_size, target_sources_num * channels_num, segment_samples)
+
+        input_dict = {'waveform': mixtures}
+        target_dict = {'waveform': targets}
+
+        return input_dict, target_dict
+
+
+class ConditionalSisoBatchDataPreprocessor:
+    def __init__(self, target_source_types: List[str]):
+        r"""Conditional single input single output (SISO) batch data
+        preprocessor. Select one target source from several target sources as
+        training target and prepare the corresponding conditional vector.
+
+        Args:
+            target_source_types: List[str], e.g., ['vocals', 'bass', ...]
+        """
+        self.target_source_types = target_source_types
+
+    def __call__(self, batch_data_dict: Dict) -> List[Dict]:
+        r"""Format waveforms and targets for training.
+
+        Args:
+            batch_data_dict: dict, e.g., {
+                'mixture': (batch_size, channels_num, segment_samples),
+                'vocals': (batch_size, channels_num, segment_samples),
+                'bass': (batch_size, channels_num, segment_samples),
+                ...,
+            }
+
+        Returns:
+            input_dict: dict, e.g., {
+                'waveform': (batch_size, channels_num, segment_samples),
+                'condition': (batch_size, target_sources_num),
+            }
+            output_dict: dict, e.g., {
+                'target': (batch_size, channels_num, segment_samples)
+            }
+        """
+
+        batch_size = len(batch_data_dict['mixture'])
+        target_sources_num = len(self.target_source_types)
+
+        assert (
+            batch_size % target_sources_num == 0
+        ), "Batch size should be \
+            evenly divided by target sources number."
+
+        mixtures = batch_data_dict['mixture']
+        # mixtures: (batch_size, channels_num, segment_samples)
+
+        conditions = torch.zeros(batch_size, target_sources_num).to(mixtures.device)
+        # conditions: (batch_size, target_sources_num)
+
+        targets = []
+
+        for n in range(batch_size):
+
+            k = n % target_sources_num  # source class index
+            source_type = self.target_source_types[k]
+
+            targets.append(batch_data_dict[source_type][n])
+
+            conditions[n, k] = 1
+
+        # conditions will looks like:
+        # [[1, 0, 0, 0],
+        #  [0, 1, 0, 0],
+        #  [0, 0, 1, 0],
+        #  [0, 0, 0, 1],
+        #  [1, 0, 0, 0],
+        #  [0, 1, 0, 0],
+        #  ...,
+        # ]
+
+        targets = torch.stack(targets, dim=0)
+        # targets: (batch_size, channels_num, segment_samples)
+
+        input_dict = {
+            'waveform': mixtures,
+            'condition': conditions,
+        }
+
+        target_dict = {'waveform': targets}
+
+        return input_dict, target_dict
+
+
+def get_batch_data_preprocessor_class(batch_data_preprocessor_type: str) -> object:
+    r"""Get batch data preprocessor class."""
+    if batch_data_preprocessor_type == 'BasicBatchDataPreprocessor':
+        return BasicBatchDataPreprocessor
+
+    elif batch_data_preprocessor_type == 'ConditionalSisoBatchDataPreprocessor':
+        return ConditionalSisoBatchDataPreprocessor
+
+    else:
+        raise NotImplementedError

bytesep/data/data_modules.py

@@ -0,0 +1,187 @@
+from typing import Dict, List, NoReturn, Optional
+
+import h5py
+import librosa
+import numpy as np
+import torch
+from pytorch_lightning.core.datamodule import LightningDataModule
+
+from bytesep.data.samplers import DistributedSamplerWrapper
+from bytesep.utils import int16_to_float32
+
+
+class DataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_sampler: object,
+        train_dataset: object,
+        num_workers: int,
+        distributed: bool,
+    ):
+        r"""Data module.
+
+        Args:
+            train_sampler: Sampler object
+            train_dataset: Dataset object
+            num_workers: int
+            distributed: bool
+        """
+        super().__init__()
+        self._train_sampler = train_sampler
+        self.train_dataset = train_dataset
+        self.num_workers = num_workers
+        self.distributed = distributed
+
+    def setup(self, stage: Optional[str] = None) -> NoReturn:
+        r"""called on every device."""
+
+        # SegmentSampler is used for selecting segments for training.
+        # On multiple devices, each SegmentSampler samples a part of mini-batch
+        # data.
+        if self.distributed:
+            self.train_sampler = DistributedSamplerWrapper(self._train_sampler)
+
+        else:
+            self.train_sampler = self._train_sampler
+
+    def train_dataloader(self) -> torch.utils.data.DataLoader:
+        r"""Get train loader."""
+        train_loader = torch.utils.data.DataLoader(
+            dataset=self.train_dataset,
+            batch_sampler=self.train_sampler,
+            collate_fn=collate_fn,
+            num_workers=self.num_workers,
+            pin_memory=True,
+        )
+
+        return train_loader
+
+
+class Dataset:
+    def __init__(self, augmentor: object, segment_samples: int):
+        r"""Used for getting data according to a meta.
+
+        Args:
+            augmentor: Augmentor class
+            segment_samples: int
+        """
+        self.augmentor = augmentor
+        self.segment_samples = segment_samples
+
+    def __getitem__(self, meta: Dict) -> Dict:
+        r"""Return data according to a meta. E.g., an input meta looks like: {
+            'vocals': [['song_A.h5', 6332760, 6465060], ['song_B.h5', 198450, 330750]],
+            'accompaniment': [['song_C.h5', 24232920, 24365250], ['song_D.h5', 1569960, 1702260]]}.
+        }
+
+        Then, vocals segments of song_A and song_B will be mixed (mix-audio augmentation).
+        Accompaniment segments of song_C and song_B will be mixed (mix-audio augmentation).
+        Finally, mixture is created by summing vocals and accompaniment.
+
+        Args:
+            meta: dict, e.g., {
+                'vocals': [['song_A.h5', 6332760, 6465060], ['song_B.h5', 198450, 330750]],
+                'accompaniment': [['song_C.h5', 24232920, 24365250], ['song_D.h5', 1569960, 1702260]]}
+            }
+
+        Returns:
+            data_dict: dict, e.g., {
+                'vocals': (channels, segments_num),
+                'accompaniment': (channels, segments_num),
+                'mixture': (channels, segments_num),
+            }
+        """
+        source_types = meta.keys()
+        data_dict = {}
+
+        for source_type in source_types:
+            # E.g., ['vocals', 'bass', ...]
+
+            waveforms = []  # Audio segments to be mix-audio augmented.
+
+            for m in meta[source_type]:
+                # E.g., {
+                #     'hdf5_path': '.../song_A.h5',
+                #     'key_in_hdf5': 'vocals',
+                #     'begin_sample': '13406400',
+                #     'end_sample': 13538700,
+                # }
+
+                hdf5_path = m['hdf5_path']
+                key_in_hdf5 = m['key_in_hdf5']
+                bgn_sample = m['begin_sample']
+                end_sample = m['end_sample']
+
+                with h5py.File(hdf5_path, 'r') as hf:
+
+                    if source_type == 'audioset':
+                        index_in_hdf5 = m['index_in_hdf5']
+                        waveform = int16_to_float32(
+                            hf['waveform'][index_in_hdf5][bgn_sample:end_sample]
+                        )
+                        waveform = waveform[None, :]
+                    else:
+                        waveform = int16_to_float32(
+                            hf[key_in_hdf5][:, bgn_sample:end_sample]
+                        )
+
+                    if self.augmentor:
+                        waveform = self.augmentor(waveform, source_type)
+
+                    waveform = librosa.util.fix_length(
+                        waveform, size=self.segment_samples, axis=1
+                    )
+                    # (channels_num, segments_num)
+
+                waveforms.append(waveform)
+            # E.g., waveforms: [(channels_num, audio_samples), (channels_num, audio_samples)]
+
+            # mix-audio augmentation
+            data_dict[source_type] = np.sum(waveforms, axis=0)
+            # data_dict[source_type]: (channels_num, audio_samples)
+
+        # data_dict looks like: {
+        #     'voclas': (channels_num, audio_samples),
+        #     'accompaniment': (channels_num, audio_samples)
+        # }
+
+        # Mix segments from different sources.
+        mixture = np.sum(
+            [data_dict[source_type] for source_type in source_types], axis=0
+        )
+        data_dict['mixture'] = mixture
+        # shape: (channels_num, audio_samples)
+
+        return data_dict
+
+
+def collate_fn(list_data_dict: List[Dict]) -> Dict:
+    r"""Collate mini-batch data to inputs and targets for training.
+
+    Args:
+        list_data_dict: e.g., [
+            {'vocals': (channels_num, segment_samples),
+             'accompaniment': (channels_num, segment_samples),
+             'mixture': (channels_num, segment_samples)
+            },
+            {'vocals': (channels_num, segment_samples),
+             'accompaniment': (channels_num, segment_samples),
+             'mixture': (channels_num, segment_samples)
+            },
+            ...]
+
+    Returns:
+        data_dict: e.g. {
+            'vocals': (batch_size, channels_num, segment_samples),
+            'accompaniment': (batch_size, channels_num, segment_samples),
+            'mixture': (batch_size, channels_num, segment_samples)
+            }
+    """
+    data_dict = {}
+
+    for key in list_data_dict[0].keys():
+        data_dict[key] = torch.Tensor(
+            np.array([data_dict[key] for data_dict in list_data_dict])
+        )
+
+    return data_dict

bytesep/data/samplers.py

@@ -0,0 +1,188 @@
+import pickle
+from typing import Dict, List, NoReturn
+
+import numpy as np
+import torch.distributed as dist
+
+
+class SegmentSampler:
+    def __init__(
+        self,
+        indexes_path: str,
+        segment_samples: int,
+        mixaudio_dict: Dict,
+        batch_size: int,
+        steps_per_epoch: int,
+        random_seed=1234,
+    ):
+        r"""Sample training indexes of sources.
+
+        Args:
+            indexes_path: str, path of indexes dict
+            segment_samplers: int
+            mixaudio_dict, dict, including hyper-parameters for mix-audio data
+                augmentation, e.g., {'voclas': 2, 'accompaniment': 2}
+            batch_size: int
+            steps_per_epoch: int, #steps_per_epoch is called an `epoch`
+            random_seed: int
+        """
+        self.segment_samples = segment_samples
+        self.mixaudio_dict = mixaudio_dict
+        self.batch_size = batch_size
+        self.steps_per_epoch = steps_per_epoch
+
+        self.meta_dict = pickle.load(open(indexes_path, "rb"))
+        # E.g., {
+        #     'vocals': [
+        #         {'hdf5_path': 'songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 0, 'end_sample': 132300}，
+        #         {'hdf5_path': 'songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4410, 'end_sample': 445410},
+        #         ...
+        #     ],
+        #     'accompaniment': [
+        #         {'hdf5_path': 'songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 0, 'end_sample': 132300}，
+        #         {'hdf5_path': 'songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4410, 'end_sample': 445410},
+        #         ...
+        #     ]
+        # }
+
+        self.source_types = self.meta_dict.keys()
+        # E.g., ['vocals', 'accompaniment']
+
+        self.pointers_dict = {source_type: 0 for source_type in self.source_types}
+        # E.g., {'vocals': 0, 'accompaniment': 0}
+
+        self.indexes_dict = {
+            source_type: np.arange(len(self.meta_dict[source_type]))
+            for source_type in self.source_types
+        }
+        # E.g. {
+        #     'vocals': [0, 1, ..., 225751],
+        #     'accompaniment': [0, 1, ..., 225751]
+        # }
+
+        self.random_state = np.random.RandomState(random_seed)
+
+        # Shuffle indexes.
+        for source_type in self.source_types:
+            self.random_state.shuffle(self.indexes_dict[source_type])
+            print("{}: {}".format(source_type, len(self.indexes_dict[source_type])))
+
+    def __iter__(self) -> List[Dict]:
+        r"""Yield a batch of meta info.
+
+        Returns:
+            batch_meta_list: (batch_size,) e.g., when mix-audio is 2, looks like [
+                {'vocals': [
+                    {'hdf5_path': 'songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 13406400, 'end_sample': 13538700},
+                    {'hdf5_path': 'songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4440870, 'end_sample': 4573170}]
+                'accompaniment': [
+                    {'hdf5_path': 'songE.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 14579460, 'end_sample': 14711760},
+                    {'hdf5_path': 'songF.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 3995460, 'end_sample': 4127760}]
+                }
+                ...
+            ]
+        """
+        batch_size = self.batch_size
+
+        while True:
+            batch_meta_dict = {source_type: [] for source_type in self.source_types}
+
+            for source_type in self.source_types:
+                # E.g., ['vocals', 'accompaniment']
+
+                # Loop until get a mini-batch.
+                while len(batch_meta_dict[source_type]) != batch_size:
+
+                    largest_index = (
+                        len(self.indexes_dict[source_type])
+                        - self.mixaudio_dict[source_type]
+                    )
+                    # E.g., 225750 = 225752 - 2
+
+                    if self.pointers_dict[source_type] > largest_index:
+
+                        # Reset pointer, and shuffle indexes.
+                        self.pointers_dict[source_type] = 0
+                        self.random_state.shuffle(self.indexes_dict[source_type])
+
+                    source_metas = []
+                    mix_audios_num = self.mixaudio_dict[source_type]
+
+                    for _ in range(mix_audios_num):
+
+                        pointer = self.pointers_dict[source_type]
+                        # E.g., 1
+
+                        index = self.indexes_dict[source_type][pointer]
+                        # E.g., 12231
+
+                        self.pointers_dict[source_type] += 1
+
+                        source_meta = self.meta_dict[source_type][index]
+                        # E.g., ['song_A.h5', 198450, 330750]
+
+                        # source_metas.append(new_source_meta)
+                        source_metas.append(source_meta)
+
+                    batch_meta_dict[source_type].append(source_metas)
+            # When mix-audio is 2, batch_meta_dict looks like: {
+            #     'vocals': [
+            #         [{'hdf5_path': 'songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 13406400, 'end_sample': 13538700},
+            #          {'hdf5_path': 'songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4440870, 'end_sample': 4573170}],
+            #         [{'hdf5_path': 'songC.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 1186290, 'end_sample': 1318590},
+            #          {'hdf5_path': 'songD.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 8462790, 'end_sample': 8595090}]
+            #     ]
+            #     'accompaniment': [
+            #         [{'hdf5_path': 'songE.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 24232950, 'end_sample': 24365250},
+            #          {'hdf5_path': 'songF.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 1569960, 'end_sample': 1702260}],
+            #         [{'hdf5_path': 'songG.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 2795940, 'end_sample': 2928240},
+            #          {'hdf5_path': 'songH.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 10923570, 'end_sample': 11055870}]
+            #     ]
+            # }
+
+            batch_meta_list = [
+                {
+                    source_type: batch_meta_dict[source_type][i]
+                    for source_type in self.source_types
+                }
+                for i in range(batch_size)
+            ]
+            # When mix-audio is 2, batch_meta_list looks like: [
+            #     {'vocals': [
+            #         {'hdf5_path': 'songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 13406400, 'end_sample': 13538700},
+            #         {'hdf5_path': 'songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4440870, 'end_sample': 4573170}]
+            #      'accompaniment': [
+            #         {'hdf5_path': 'songE.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 14579460, 'end_sample': 14711760},
+            #         {'hdf5_path': 'songF.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 3995460, 'end_sample': 4127760}]
+            #     }
+            #     ...
+            # ]
+
+            yield batch_meta_list
+
+    def __len__(self) -> int:
+        return self.steps_per_epoch
+
+    def state_dict(self) -> Dict:
+        state = {'pointers_dict': self.pointers_dict, 'indexes_dict': self.indexes_dict}
+        return state
+
+    def load_state_dict(self, state) -> NoReturn:
+        self.pointers_dict = state['pointers_dict']
+        self.indexes_dict = state['indexes_dict']
+
+
+class DistributedSamplerWrapper:
+    def __init__(self, sampler):
+        r"""Distributed wrapper of sampler."""
+        self.sampler = sampler
+
+    def __iter__(self):
+        num_replicas = dist.get_world_size()
+        rank = dist.get_rank()
+
+        for indices in self.sampler:
+            yield indices[rank::num_replicas]
+
+    def __len__(self) -> int:
+        return len(self.sampler)

bytesep/dataset_creation/create_evaluation_audios/piano-symphony.py

@@ -0,0 +1,160 @@
+import argparse
+import os
+from typing import NoReturn
+
+import librosa
+import numpy as np
+import soundfile
+
+from bytesep.dataset_creation.pack_audios_to_hdf5s.instruments_solo import (
+    read_csv as read_instruments_solo_csv,
+)
+from bytesep.dataset_creation.pack_audios_to_hdf5s.maestro import (
+    read_csv as read_maestro_csv,
+)
+from bytesep.utils import load_random_segment
+
+
+def create_evaluation(args) -> NoReturn:
+    r"""Random mix and write out audios for evaluation.
+
+    Args:
+        piano_dataset_dir: str, the directory of the piano dataset
+        symphony_dataset_dir: str, the directory of the symphony dataset
+        evaluation_audios_dir: str, the directory to write out randomly selected and mixed audio segments
+        sample_rate: int
+        channels: int, e.g., 1 | 2
+        evaluation_segments_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    piano_dataset_dir = args.piano_dataset_dir
+    symphony_dataset_dir = args.symphony_dataset_dir
+    evaluation_audios_dir = args.evaluation_audios_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    evaluation_segments_num = args.evaluation_segments_num
+    mono = True if channels == 1 else False
+
+    split = 'test'
+    segment_seconds = 10.0
+
+    random_state = np.random.RandomState(1234)
+
+    piano_meta_csv = os.path.join(piano_dataset_dir, 'maestro-v2.0.0.csv')
+    piano_names_dict = read_maestro_csv(piano_meta_csv)
+    piano_audio_names = piano_names_dict[split]
+
+    symphony_meta_csv = os.path.join(symphony_dataset_dir, 'validation.csv')
+    symphony_names_dict = read_instruments_solo_csv(symphony_meta_csv)
+    symphony_audio_names = symphony_names_dict[split]
+
+    for source_type in ['piano', 'symphony', 'mixture']:
+        output_dir = os.path.join(evaluation_audios_dir, split, source_type)
+        os.makedirs(output_dir, exist_ok=True)
+
+    for n in range(evaluation_segments_num):
+
+        print('{} / {}'.format(n, evaluation_segments_num))
+
+        # Randomly select and write out a clean piano segment.
+        piano_audio_name = random_state.choice(piano_audio_names)
+        piano_audio_path = os.path.join(piano_dataset_dir, piano_audio_name)
+
+        piano_audio = load_random_segment(
+            audio_path=piano_audio_path,
+            random_state=random_state,
+            segment_seconds=segment_seconds,
+            mono=mono,
+            sample_rate=sample_rate,
+        )
+
+        output_piano_path = os.path.join(
+            evaluation_audios_dir, split, 'piano', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_piano_path, data=piano_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_piano_path))
+
+        # Randomly select and write out a clean symphony segment.
+        symphony_audio_name = random_state.choice(symphony_audio_names)
+        symphony_audio_path = os.path.join(
+            symphony_dataset_dir, "mp3s", symphony_audio_name
+        )
+
+        symphony_audio = load_random_segment(
+            audio_path=symphony_audio_path,
+            random_state=random_state,
+            segment_seconds=segment_seconds,
+            mono=mono,
+            sample_rate=sample_rate,
+        )
+
+        output_symphony_path = os.path.join(
+            evaluation_audios_dir, split, 'symphony', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_symphony_path, data=symphony_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_symphony_path))
+
+        # Mix piano and symphony segments and write out a mixture segment.
+        mixture_audio = symphony_audio + piano_audio
+        output_mixture_path = os.path.join(
+            evaluation_audios_dir, split, 'mixture', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_mixture_path, data=mixture_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_mixture_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--piano_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the piano dataset.",
+    )
+    parser.add_argument(
+        "--symphony_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the symphony dataset.",
+    )
+    parser.add_argument(
+        "--evaluation_audios_dir",
+        type=str,
+        required=True,
+        help="The directory to write out randomly selected and mixed audio segments.",
+    )
+    parser.add_argument(
+        "--sample_rate",
+        type=int,
+        required=True,
+        help="Sample rate.",
+    )
+    parser.add_argument(
+        "--channels",
+        type=int,
+        required=True,
+        help="Audio channels, e.g, 1 or 2.",
+    )
+    parser.add_argument(
+        "--evaluation_segments_num",
+        type=int,
+        required=True,
+        help="The number of segments to create for evaluation.",
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    create_evaluation(args)

bytesep/dataset_creation/create_evaluation_audios/vctk-musdb18.py

@@ -0,0 +1,164 @@
+import argparse
+import os
+import soundfile
+from typing import NoReturn
+
+import musdb
+import numpy as np
+
+from bytesep.utils import load_audio
+
+
+def create_evaluation(args) -> NoReturn:
+    r"""Random mix and write out audios for evaluation.
+
+    Args:
+        vctk_dataset_dir: str, the directory of the VCTK dataset
+        symphony_dataset_dir: str, the directory of the symphony dataset
+        evaluation_audios_dir: str, the directory to write out randomly selected and mixed audio segments
+        sample_rate: int
+        channels: int, e.g., 1 | 2
+        evaluation_segments_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    vctk_dataset_dir = args.vctk_dataset_dir
+    musdb18_dataset_dir = args.musdb18_dataset_dir
+    evaluation_audios_dir = args.evaluation_audios_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    evaluation_segments_num = args.evaluation_segments_num
+    mono = True if channels == 1 else False
+
+    split = 'test'
+    random_state = np.random.RandomState(1234)
+
+    # paths
+    audios_dir = os.path.join(vctk_dataset_dir, "wav48", split)
+
+    for source_type in ['speech', 'music', 'mixture']:
+        output_dir = os.path.join(evaluation_audios_dir, split, source_type)
+        os.makedirs(output_dir, exist_ok=True)
+
+    # Get VCTK audio paths.
+    speech_audio_paths = []
+    speaker_ids = sorted(os.listdir(audios_dir))
+
+    for speaker_id in speaker_ids:
+        speaker_audios_dir = os.path.join(audios_dir, speaker_id)
+
+        audio_names = sorted(os.listdir(speaker_audios_dir))
+
+        for audio_name in audio_names:
+            speaker_audio_path = os.path.join(speaker_audios_dir, audio_name)
+            speech_audio_paths.append(speaker_audio_path)
+
+    # Get Musdb18 audio paths.
+    mus = musdb.DB(root=musdb18_dataset_dir, subsets=[split])
+    track_indexes = np.arange(len(mus.tracks))
+
+    for n in range(evaluation_segments_num):
+
+        print('{} / {}'.format(n, evaluation_segments_num))
+
+        # Randomly select and write out a clean speech segment.
+        speech_audio_path = random_state.choice(speech_audio_paths)
+
+        speech_audio = load_audio(
+            audio_path=speech_audio_path, mono=mono, sample_rate=sample_rate
+        )
+        # (channels_num, audio_samples)
+
+        if channels == 2:
+            speech_audio = np.tile(speech_audio, (2, 1))
+        # (channels_num, audio_samples)
+
+        output_speech_path = os.path.join(
+            evaluation_audios_dir, split, 'speech', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_speech_path, data=speech_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_speech_path))
+
+        # Randomly select and write out a clean music segment.
+        track_index = random_state.choice(track_indexes)
+        track = mus[track_index]
+
+        segment_samples = speech_audio.shape[1]
+        start_sample = int(
+            random_state.uniform(0.0, segment_samples - speech_audio.shape[1])
+        )
+
+        music_audio = track.audio[start_sample : start_sample + segment_samples, :].T
+        # (channels_num, audio_samples)
+
+        output_music_path = os.path.join(
+            evaluation_audios_dir, split, 'music', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_music_path, data=music_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_music_path))
+
+        # Mix speech and music segments and write out a mixture segment.
+        mixture_audio = speech_audio + music_audio
+        # (channels_num, audio_samples)
+
+        output_mixture_path = os.path.join(
+            evaluation_audios_dir, split, 'mixture', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_mixture_path, data=mixture_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_mixture_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--vctk_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the VCTK dataset.",
+    )
+    parser.add_argument(
+        "--musdb18_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the MUSDB18 dataset.",
+    )
+    parser.add_argument(
+        "--evaluation_audios_dir",
+        type=str,
+        required=True,
+        help="The directory to write out randomly selected and mixed audio segments.",
+    )
+    parser.add_argument(
+        "--sample_rate",
+        type=int,
+        required=True,
+        help="Sample rate",
+    )
+    parser.add_argument(
+        "--channels",
+        type=int,
+        required=True,
+        help="Audio channels, e.g, 1 or 2.",
+    )
+    parser.add_argument(
+        "--evaluation_segments_num",
+        type=int,
+        required=True,
+        help="The number of segments to create for evaluation.",
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    create_evaluation(args)

bytesep/dataset_creation/create_evaluation_audios/violin-piano.py

@@ -0,0 +1,162 @@
+import argparse
+import os
+from typing import NoReturn
+
+import librosa
+import numpy as np
+import soundfile
+
+from bytesep.dataset_creation.pack_audios_to_hdf5s.instruments_solo import (
+    read_csv as read_instruments_solo_csv,
+)
+from bytesep.dataset_creation.pack_audios_to_hdf5s.maestro import (
+    read_csv as read_maestro_csv,
+)
+from bytesep.utils import load_random_segment
+
+
+def create_evaluation(args) -> NoReturn:
+    r"""Random mix and write out audios for evaluation.
+
+    Args:
+        violin_dataset_dir: str, the directory of the violin dataset
+        piano_dataset_dir: str, the directory of the piano dataset
+        evaluation_audios_dir: str, the directory to write out randomly selected and mixed audio segments
+        sample_rate: int
+        channels: int, e.g., 1 | 2
+        evaluation_segments_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    violin_dataset_dir = args.violin_dataset_dir
+    piano_dataset_dir = args.piano_dataset_dir
+    evaluation_audios_dir = args.evaluation_audios_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    evaluation_segments_num = args.evaluation_segments_num
+    mono = True if channels == 1 else False
+
+    split = 'test'
+    segment_seconds = 10.0
+
+    random_state = np.random.RandomState(1234)
+
+    violin_meta_csv = os.path.join(violin_dataset_dir, 'validation.csv')
+    violin_names_dict = read_instruments_solo_csv(violin_meta_csv)
+    violin_audio_names = violin_names_dict['{}'.format(split)]
+
+    piano_meta_csv = os.path.join(piano_dataset_dir, 'maestro-v2.0.0.csv')
+    piano_names_dict = read_maestro_csv(piano_meta_csv)
+    piano_audio_names = piano_names_dict['{}'.format(split)]
+
+    for source_type in ['violin', 'piano', 'mixture']:
+        output_dir = os.path.join(evaluation_audios_dir, split, source_type)
+        os.makedirs(output_dir, exist_ok=True)
+
+    for n in range(evaluation_segments_num):
+
+        print('{} / {}'.format(n, evaluation_segments_num))
+
+        # Randomly select and write out a clean violin segment.
+        violin_audio_name = random_state.choice(violin_audio_names)
+        violin_audio_path = os.path.join(violin_dataset_dir, "mp3s", violin_audio_name)
+
+        violin_audio = load_random_segment(
+            audio_path=violin_audio_path,
+            random_state=random_state,
+            segment_seconds=segment_seconds,
+            mono=mono,
+            sample_rate=sample_rate,
+        )
+        # (channels_num, audio_samples)
+
+        output_violin_path = os.path.join(
+            evaluation_audios_dir, split, 'violin', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_violin_path, data=violin_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_violin_path))
+
+        # Randomly select and write out a clean piano segment.
+        piano_audio_name = random_state.choice(piano_audio_names)
+        piano_audio_path = os.path.join(piano_dataset_dir, piano_audio_name)
+
+        piano_audio = load_random_segment(
+            audio_path=piano_audio_path,
+            random_state=random_state,
+            segment_seconds=segment_seconds,
+            mono=mono,
+            sample_rate=sample_rate,
+        )
+        # (channels_num, audio_samples)
+
+        output_piano_path = os.path.join(
+            evaluation_audios_dir, split, 'piano', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_piano_path, data=piano_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_piano_path))
+
+        # Mix violin and piano segments and write out a mixture segment.
+        mixture_audio = violin_audio + piano_audio
+        # (channels_num, audio_samples)
+
+        output_mixture_path = os.path.join(
+            evaluation_audios_dir, split, 'mixture', '{:04d}.wav'.format(n)
+        )
+        soundfile.write(
+            file=output_mixture_path, data=mixture_audio.T, samplerate=sample_rate
+        )
+        print("Write out to {}".format(output_mixture_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--violin_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the violin dataset.",
+    )
+    parser.add_argument(
+        "--piano_dataset_dir",
+        type=str,
+        required=True,
+        help="The directory of the piano dataset.",
+    )
+    parser.add_argument(
+        "--evaluation_audios_dir",
+        type=str,
+        required=True,
+        help="The directory to write out randomly selected and mixed audio segments.",
+    )
+    parser.add_argument(
+        "--sample_rate",
+        type=int,
+        required=True,
+        help="Sample rate",
+    )
+    parser.add_argument(
+        "--channels",
+        type=int,
+        required=True,
+        help="Audio channels, e.g, 1 or 2.",
+    )
+    parser.add_argument(
+        "--evaluation_segments_num",
+        type=int,
+        required=True,
+        help="The number of segments to create for evaluation.",
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    create_evaluation(args)

bytesep/dataset_creation/create_indexes/create_indexes.py

@@ -0,0 +1,142 @@
+import argparse
+import os
+import pickle
+from typing import NoReturn
+
+import h5py
+
+from bytesep.utils import read_yaml
+
+
+def create_indexes(args) -> NoReturn:
+    r"""Create and write out training indexes into disk. The indexes may contain
+    information from multiple datasets. During training, training indexes will
+    be shuffled and iterated for selecting segments to be mixed. E.g., the
+    training indexes_dict looks like: {
+        'vocals': [
+            {'hdf5_path': '.../songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 0, 'end_sample': 132300}
+            {'hdf5_path': '.../songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4410, 'end_sample': 136710}
+            ...
+        ]
+        'accompaniment': [
+            {'hdf5_path': '.../songA.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 0, 'end_sample': 132300}
+            {'hdf5_path': '.../songB.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 4410, 'end_sample': 136710}
+            ...
+        ]
+    }
+    """
+
+    # Arugments & parameters
+    workspace = args.workspace
+    config_yaml = args.config_yaml
+
+    # Only create indexes for training, because evalution is on entire pieces.
+    split = "train"
+
+    # Read config file.
+    configs = read_yaml(config_yaml)
+
+    sample_rate = configs["sample_rate"]
+    segment_samples = int(configs["segment_seconds"] * sample_rate)
+
+    # Path to write out index.
+    indexes_path = os.path.join(workspace, configs[split]["indexes"])
+    os.makedirs(os.path.dirname(indexes_path), exist_ok=True)
+
+    source_types = configs[split]["source_types"].keys()
+    # E.g., ['vocals', 'accompaniment']
+
+    indexes_dict = {source_type: [] for source_type in source_types}
+    # E.g., indexes_dict will looks like: {
+    #     'vocals': [
+    #         {'hdf5_path': '.../songA.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 0, 'end_sample': 132300}
+    #         {'hdf5_path': '.../songB.h5', 'key_in_hdf5': 'vocals', 'begin_sample': 4410, 'end_sample': 136710}
+    #         ...
+    #     ]
+    #     'accompaniment': [
+    #         {'hdf5_path': '.../songA.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 0, 'end_sample': 132300}
+    #         {'hdf5_path': '.../songB.h5', 'key_in_hdf5': 'accompaniment', 'begin_sample': 4410, 'end_sample': 136710}
+    #         ...
+    #     ]
+    # }
+
+    # Get training indexes for each source type.
+    for source_type in source_types:
+        # E.g., ['vocals', 'bass', ...]
+
+        print("--- {} ---".format(source_type))
+
+        dataset_types = configs[split]["source_types"][source_type]
+        # E.g., ['musdb18', ...]
+
+        # Each source can come from mulitple datasets.
+        for dataset_type in dataset_types:
+
+            hdf5s_dir = os.path.join(
+                workspace, dataset_types[dataset_type]["hdf5s_directory"]
+            )
+
+            hop_samples = int(dataset_types[dataset_type]["hop_seconds"] * sample_rate)
+
+            key_in_hdf5 = dataset_types[dataset_type]["key_in_hdf5"]
+            # E.g., 'vocals'
+
+            hdf5_names = sorted(os.listdir(hdf5s_dir))
+            print("Hdf5 files num: {}".format(len(hdf5_names)))
+
+            # Traverse all packed hdf5 files of a dataset.
+            for n, hdf5_name in enumerate(hdf5_names):
+
+                print(n, hdf5_name)
+                hdf5_path = os.path.join(hdf5s_dir, hdf5_name)
+
+                with h5py.File(hdf5_path, "r") as hf:
+
+                    bgn_sample = 0
+                    while bgn_sample + segment_samples < hf[key_in_hdf5].shape[-1]:
+                        meta = {
+                            'hdf5_path': hdf5_path,
+                            'key_in_hdf5': key_in_hdf5,
+                            'begin_sample': bgn_sample,
+                            'end_sample': bgn_sample + segment_samples,
+                        }
+                        indexes_dict[source_type].append(meta)
+
+                        bgn_sample += hop_samples
+
+                    # If the audio length is shorter than the segment length,
+                    # then use the entire audio as a segment.
+                    if bgn_sample == 0:
+                        meta = {
+                            'hdf5_path': hdf5_path,
+                            'key_in_hdf5': key_in_hdf5,
+                            'begin_sample': 0,
+                            'end_sample': segment_samples,
+                        }
+                        indexes_dict[source_type].append(meta)
+
+        print(
+            "Total indexes for {}: {}".format(
+                source_type, len(indexes_dict[source_type])
+            )
+        )
+
+    pickle.dump(indexes_dict, open(indexes_path, "wb"))
+    print("Write index dict to {}".format(indexes_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--workspace", type=str, required=True, help="Directory of workspace."
+    )
+    parser.add_argument(
+        "--config_yaml", type=str, required=True, help="User defined config file."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Create training indexes.
+    create_indexes(args)

bytesep/dataset_creation/pack_audios_to_hdf5s/instruments_solo.py

@@ -0,0 +1,173 @@
+import argparse
+import os
+import pathlib
+import time
+from concurrent.futures import ProcessPoolExecutor
+from typing import Dict, List, NoReturn
+
+import h5py
+import numpy as np
+import pandas as pd
+
+from bytesep.utils import float32_to_int16, load_audio
+
+
+def read_csv(meta_csv) -> Dict:
+    r"""Get train & test names from csv.
+
+    Args:
+        meta_csv: str
+
+    Returns:
+        names_dict: dict, e.g., {
+            'train', ['songA.mp3', 'songB.mp3', ...],
+            'test': ['songE.mp3', 'songF.mp3', ...]
+        }
+    """
+    df = pd.read_csv(meta_csv, sep=',')
+
+    names_dict = {}
+
+    for split in ['train', 'test']:
+        audio_indexes = df['split'] == split
+        audio_names = list(df['audio_name'][audio_indexes])
+        audio_names = [
+            '{}.mp3'.format(pathlib.Path(audio_name).stem) for audio_name in audio_names
+        ]
+        names_dict[split] = audio_names
+
+    return names_dict
+
+
+def pack_audios_to_hdf5s(args) -> NoReturn:
+    r"""Pack (resampled) audio files into hdf5 files to speed up loading.
+
+    Args:
+        dataset_dir: str
+        split: str, 'train' | 'test'
+        source_type: str
+        hdf5s_dir: str, directory to write out hdf5 files
+        sample_rate: int
+        channels_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    dataset_dir = args.dataset_dir
+    split = args.split
+    source_type = args.source_type
+    hdf5s_dir = args.hdf5s_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    mono = True if channels == 1 else False
+
+    # Only pack data for training data.
+    assert split == "train"
+
+    # paths
+    audios_dir = os.path.join(dataset_dir, 'mp3s')
+    meta_csv = os.path.join(dataset_dir, 'validation.csv')
+
+    os.makedirs(hdf5s_dir, exist_ok=True)
+
+    # Read train & test names.
+    names_dict = read_csv(meta_csv)
+
+    audio_names = names_dict[split]
+
+    params = []
+
+    for audio_index, audio_name in enumerate(audio_names):
+
+        audio_path = os.path.join(audios_dir, audio_name)
+
+        hdf5_path = os.path.join(
+            hdf5s_dir, "{}.h5".format(pathlib.Path(audio_name).stem)
+        )
+
+        param = (
+            audio_index,
+            audio_name,
+            source_type,
+            audio_path,
+            mono,
+            sample_rate,
+            hdf5_path,
+        )
+        params.append(param)
+
+    # Uncomment for debug.
+    # write_single_audio_to_hdf5(params[0])
+    # os._exit()
+
+    pack_hdf5s_time = time.time()
+
+    with ProcessPoolExecutor(max_workers=None) as pool:
+        # Maximum works on the machine
+        pool.map(write_single_audio_to_hdf5, params)
+
+    print("Pack hdf5 time: {:.3f} s".format(time.time() - pack_hdf5s_time))
+
+
+def write_single_audio_to_hdf5(param: List) -> NoReturn:
+    r"""Write single audio into hdf5 file."""
+
+    (
+        audio_index,
+        audio_name,
+        source_type,
+        audio_path,
+        mono,
+        sample_rate,
+        hdf5_path,
+    ) = param
+
+    with h5py.File(hdf5_path, "w") as hf:
+
+        hf.attrs.create("audio_name", data=audio_name, dtype="S100")
+        hf.attrs.create("sample_rate", data=sample_rate, dtype=np.int32)
+
+        audio = load_audio(audio_path=audio_path, mono=mono, sample_rate=sample_rate)
+        # audio: (channels_num, audio_samples)
+
+        hf.create_dataset(
+            name=source_type, data=float32_to_int16(audio), dtype=np.int16
+        )
+
+    print('{} Write hdf5 to {}'.format(audio_index, hdf5_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory of the instruments solo dataset.",
+    )
+    parser.add_argument("--split", type=str, required=True, choices=["train", "test"])
+    parser.add_argument(
+        "--source_type",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--hdf5s_dir",
+        type=str,
+        required=True,
+        help="Directory to write out hdf5 files.",
+    )
+    parser.add_argument("--sample_rate", type=int, required=True, help="Sample rate.")
+    parser.add_argument(
+        "--channels", type=int, required=True, help="Use 1 for mono, 2 for stereo."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Pack audios to hdf5 files.
+    pack_audios_to_hdf5s(args)

bytesep/dataset_creation/pack_audios_to_hdf5s/maestro.py

@@ -0,0 +1,136 @@
+import argparse
+import os
+import pathlib
+import time
+from concurrent.futures import ProcessPoolExecutor
+from typing import Dict, NoReturn
+
+import pandas as pd
+
+from bytesep.dataset_creation.pack_audios_to_hdf5s.instruments_solo import (
+    write_single_audio_to_hdf5,
+)
+
+
+def read_csv(meta_csv) -> Dict:
+    r"""Get train & test names from csv.
+
+    Args:
+        meta_csv: str
+
+    Returns:
+        names_dict: dict, e.g., {
+            'train', ['a1.mp3', 'a2.mp3'],
+            'test': ['b1.mp3', 'b2.mp3']
+        }
+    """
+    df = pd.read_csv(meta_csv, sep=',')
+
+    names_dict = {}
+
+    for split in ['train', 'test']:
+        audio_indexes = df['split'] == split
+        audio_names = list(df['audio_filename'][audio_indexes])
+        names_dict[split] = audio_names
+
+    return names_dict
+
+
+def pack_audios_to_hdf5s(args) -> NoReturn:
+    r"""Pack (resampled) audio files into hdf5 files to speed up loading.
+
+    Args:
+        dataset_dir: str
+        split: str, 'train' | 'test'
+        hdf5s_dir: str, directory to write out hdf5 files
+        sample_rate: int
+        channels_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    dataset_dir = args.dataset_dir
+    split = args.split
+    hdf5s_dir = args.hdf5s_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    mono = True if channels == 1 else False
+
+    source_type = "piano"
+
+    # Only pack data for training data.
+    assert split == "train"
+
+    # paths
+    meta_csv = os.path.join(dataset_dir, 'maestro-v2.0.0.csv')
+
+    os.makedirs(hdf5s_dir, exist_ok=True)
+
+    # Read train & test names.
+    names_dict = read_csv(meta_csv)
+
+    audio_names = names_dict['{}'.format(split)]
+
+    params = []
+
+    for audio_index, audio_name in enumerate(audio_names):
+
+        audio_path = os.path.join(dataset_dir, audio_name)
+
+        hdf5_path = os.path.join(
+            hdf5s_dir, "{}.h5".format(pathlib.Path(audio_name).stem)
+        )
+
+        param = (
+            audio_index,
+            audio_name,
+            source_type,
+            audio_path,
+            mono,
+            sample_rate,
+            hdf5_path,
+        )
+        params.append(param)
+
+    # Uncomment for debug.
+    # write_single_audio_to_hdf5(params[0])
+    # os._exit(0)
+
+    pack_hdf5s_time = time.time()
+
+    with ProcessPoolExecutor(max_workers=None) as pool:
+        # Maximum works on the machine
+        pool.map(write_single_audio_to_hdf5, params)
+
+    print("Pack hdf5 time: {:.3f} s".format(time.time() - pack_hdf5s_time))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory of the MAESTRO dataset.",
+    )
+    parser.add_argument("--split", type=str, required=True, choices=["train", "test"])
+    parser.add_argument(
+        "--hdf5s_dir",
+        type=str,
+        required=True,
+        help="Directory to write out hdf5 files.",
+    )
+    parser.add_argument("--sample_rate", type=int, required=True, help="Sample rate.")
+    parser.add_argument(
+        "--channels", type=int, required=True, help="Use 1 for mono, 2 for stereo."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Pack audios to hdf5 files.
+    pack_audios_to_hdf5s(args)

bytesep/dataset_creation/pack_audios_to_hdf5s/musdb18.py

@@ -0,0 +1,207 @@
+import argparse
+import os
+import time
+from concurrent.futures import ProcessPoolExecutor
+from typing import NoReturn
+
+import h5py
+import librosa
+import musdb
+import numpy as np
+
+from bytesep.utils import float32_to_int16
+
+# Source types of the MUSDB18 dataset.
+SOURCE_TYPES = ["vocals", "drums", "bass", "other", "accompaniment"]
+
+
+def pack_audios_to_hdf5s(args) -> NoReturn:
+    r"""Pack (resampled) audio files into hdf5 files to speed up loading.
+
+    Args:
+        dataset_dir: str
+        subset: str, 'train' | 'test'
+        split: str, '' | 'train' | 'valid'
+        hdf5s_dir: str, directory to write out hdf5 files
+        sample_rate: int
+        channels_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    dataset_dir = args.dataset_dir
+    subset = args.subset
+    split = None if args.split == "" else args.split
+    hdf5s_dir = args.hdf5s_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+
+    mono = True if channels == 1 else False
+    source_types = SOURCE_TYPES
+    resample_type = "kaiser_fast"
+
+    # Paths
+    os.makedirs(hdf5s_dir, exist_ok=True)
+
+    # Dataset of corresponding subset and split.
+    mus = musdb.DB(root=dataset_dir, subsets=[subset], split=split)
+    print("Subset: {}, Split: {}, Total pieces: {}".format(subset, split, len(mus)))
+
+    params = []  # A list of params for multiple processing.
+
+    for track_index in range(len(mus.tracks)):
+
+        param = (
+            dataset_dir,
+            subset,
+            split,
+            track_index,
+            source_types,
+            mono,
+            sample_rate,
+            resample_type,
+            hdf5s_dir,
+        )
+
+        params.append(param)
+
+    # Uncomment for debug.
+    # write_single_audio_to_hdf5(params[0])
+    # os._exit(0)
+
+    pack_hdf5s_time = time.time()
+
+    with ProcessPoolExecutor(max_workers=None) as pool:
+        # Maximum works on the machine
+        pool.map(write_single_audio_to_hdf5, params)
+
+    print("Pack hdf5 time: {:.3f} s".format(time.time() - pack_hdf5s_time))
+
+
+def write_single_audio_to_hdf5(param) -> NoReturn:
+    r"""Write single audio into hdf5 file."""
+    (
+        dataset_dir,
+        subset,
+        split,
+        track_index,
+        source_types,
+        mono,
+        sample_rate,
+        resample_type,
+        hdf5s_dir,
+    ) = param
+
+    # Dataset of corresponding subset and split.
+    mus = musdb.DB(root=dataset_dir, subsets=[subset], split=split)
+    track = mus.tracks[track_index]
+
+    # Path to write out hdf5 file.
+    hdf5_path = os.path.join(hdf5s_dir, "{}.h5".format(track.name))
+
+    with h5py.File(hdf5_path, "w") as hf:
+
+        hf.attrs.create("audio_name", data=track.name.encode(), dtype="S100")
+        hf.attrs.create("sample_rate", data=sample_rate, dtype=np.int32)
+
+        for source_type in source_types:
+
+            audio = track.targets[source_type].audio.T
+            # (channels_num, audio_samples)
+
+            # Preprocess audio to mono / stereo, and resample.
+            audio = preprocess_audio(
+                audio, mono, track.rate, sample_rate, resample_type
+            )
+            # audio = load_audio(audio_path=audio_path, mono=mono, sample_rate=sample_rate)
+            # (channels_num, audio_samples) | (audio_samples,)
+
+            hf.create_dataset(
+                name=source_type, data=float32_to_int16(audio), dtype=np.int16
+            )
+
+        # Mixture
+        audio = track.audio.T
+        # (channels_num, audio_samples)
+
+        # Preprocess audio to mono / stereo, and resample.
+        audio = preprocess_audio(audio, mono, track.rate, sample_rate, resample_type)
+        # (channels_num, audio_samples)
+
+        hf.create_dataset(name="mixture", data=float32_to_int16(audio), dtype=np.int16)
+
+    print("{} Write to {}, {}".format(track_index, hdf5_path, audio.shape))
+
+
+def preprocess_audio(audio, mono, origin_sr, sr, resample_type) -> np.array:
+    r"""Preprocess audio to mono / stereo, and resample.
+
+    Args:
+        audio: (channels_num, audio_samples), input audio
+        mono: bool
+        origin_sr: float, original sample rate
+        sr: float, target sample rate
+        resample_type: str, e.g., 'kaiser_fast'
+
+    Returns:
+        output: ndarray, output audio
+    """
+    if mono:
+        audio = np.mean(audio, axis=0)
+        # (audio_samples,)
+
+    output = librosa.core.resample(
+        audio, orig_sr=origin_sr, target_sr=sr, res_type=resample_type
+    )
+    # (audio_samples,) | (channels_num, audio_samples)
+
+    if output.ndim == 1:
+        output = output[None, :]
+        # (1, audio_samples,)
+
+    return output
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory of the MUSDB18 dataset.",
+    )
+    parser.add_argument(
+        "--subset",
+        type=str,
+        required=True,
+        choices=["train", "test"],
+        help="Train subset: 100 pieces; test subset: 50 pieces.",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        required=True,
+        choices=["", "train", "valid"],
+        help="Use '' to use all 100 pieces to train. Use 'train' to use 86 \
+            pieces for train, and use 'test' to use 14 pieces for valid.",
+    )
+    parser.add_argument(
+        "--hdf5s_dir",
+        type=str,
+        required=True,
+        help="Directory to write out hdf5 files.",
+    )
+    parser.add_argument("--sample_rate", type=int, required=True, help="Sample rate.")
+    parser.add_argument(
+        "--channels", type=int, required=True, help="Use 1 for mono, 2 for stereo."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Pack audios into hdf5 files.
+    pack_audios_to_hdf5s(args)

bytesep/dataset_creation/pack_audios_to_hdf5s/vctk.py

@@ -0,0 +1,114 @@
+import argparse
+import os
+import pathlib
+import time
+from concurrent.futures import ProcessPoolExecutor
+from typing import NoReturn
+
+from bytesep.dataset_creation.pack_audios_to_hdf5s.instruments_solo import (
+    write_single_audio_to_hdf5,
+)
+
+
+def pack_audios_to_hdf5s(args) -> NoReturn:
+    r"""Pack (resampled) audio files into hdf5 files to speed up loading.
+
+    Args:
+        dataset_dir: str
+        split: str, 'train' | 'test'
+        hdf5s_dir: str, directory to write out hdf5 files
+        sample_rate: int
+        channels_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    dataset_dir = args.dataset_dir
+    split = args.split
+    hdf5s_dir = args.hdf5s_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    mono = True if channels == 1 else False
+
+    source_type = "speech"
+
+    # Only pack data for training data.
+    assert split == "train"
+
+    audios_dir = os.path.join(dataset_dir, 'wav48', split)
+    os.makedirs(hdf5s_dir, exist_ok=True)
+
+    speaker_ids = sorted(os.listdir(audios_dir))
+
+    params = []
+    audio_index = 0
+
+    for speaker_id in speaker_ids:
+
+        speaker_audios_dir = os.path.join(audios_dir, speaker_id)
+
+        audio_names = sorted(os.listdir(speaker_audios_dir))
+
+        for audio_name in audio_names:
+
+            audio_path = os.path.join(speaker_audios_dir, audio_name)
+
+            hdf5_path = os.path.join(
+                hdf5s_dir, "{}.h5".format(pathlib.Path(audio_name).stem)
+            )
+
+            param = (
+                audio_index,
+                audio_name,
+                source_type,
+                audio_path,
+                mono,
+                sample_rate,
+                hdf5_path,
+            )
+            params.append(param)
+
+            audio_index += 1
+
+    # Uncomment for debug.
+    # write_single_audio_to_hdf5(params[0])
+    # os._exit(0)
+
+    pack_hdf5s_time = time.time()
+
+    with ProcessPoolExecutor(max_workers=None) as pool:
+        # Maximum works on the machine
+        pool.map(write_single_audio_to_hdf5, params)
+
+    print("Pack hdf5 time: {:.3f} s".format(time.time() - pack_hdf5s_time))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory of the VCTK dataset.",
+    )
+    parser.add_argument("--split", type=str, required=True, choices=["train", "test"])
+    parser.add_argument(
+        "--hdf5s_dir",
+        type=str,
+        required=True,
+        help="Directory to write out hdf5 files.",
+    )
+    parser.add_argument("--sample_rate", type=int, required=True, help="Sample rate.")
+    parser.add_argument(
+        "--channels", type=int, required=True, help="Use 1 for mono, 2 for stereo."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Pack audios into hdf5 files.
+    pack_audios_to_hdf5s(args)

bytesep/dataset_creation/pack_audios_to_hdf5s/voicebank-demand.py

@@ -0,0 +1,143 @@
+import argparse
+import os
+import pathlib
+import time
+from concurrent.futures import ProcessPoolExecutor
+from typing import List, NoReturn
+
+import h5py
+import numpy as np
+
+from bytesep.utils import float32_to_int16, load_audio
+
+
+def pack_audios_to_hdf5s(args) -> NoReturn:
+    r"""Pack (resampled) audio files into hdf5 files to speed up loading.
+
+    Args:
+        dataset_dir: str
+        split: str, 'train' | 'test'
+        hdf5s_dir: str, directory to write out hdf5 files
+        sample_rate: int
+        channels_num: int
+        mono: bool
+
+    Returns:
+        NoReturn
+    """
+
+    # arguments & parameters
+    dataset_dir = args.dataset_dir
+    split = args.split
+    hdf5s_dir = args.hdf5s_dir
+    sample_rate = args.sample_rate
+    channels = args.channels
+    mono = True if channels == 1 else False
+
+    # Only pack data for training data.
+    assert split == "train"
+
+    speech_dir = os.path.join(dataset_dir, "clean_{}set_wav".format(split))
+    mixture_dir = os.path.join(dataset_dir, "noisy_{}set_wav".format(split))
+
+    os.makedirs(hdf5s_dir, exist_ok=True)
+
+    # Read names.
+    audio_names = sorted(os.listdir(speech_dir))
+
+    params = []
+
+    for audio_index, audio_name in enumerate(audio_names):
+
+        speech_path = os.path.join(speech_dir, audio_name)
+        mixture_path = os.path.join(mixture_dir, audio_name)
+
+        hdf5_path = os.path.join(
+            hdf5s_dir, "{}.h5".format(pathlib.Path(audio_name).stem)
+        )
+
+        param = (
+            audio_index,
+            audio_name,
+            speech_path,
+            mixture_path,
+            mono,
+            sample_rate,
+            hdf5_path,
+        )
+        params.append(param)
+
+    # Uncomment for debug.
+    # write_single_audio_to_hdf5(params[0])
+    # os._exit(0)
+
+    pack_hdf5s_time = time.time()
+
+    with ProcessPoolExecutor(max_workers=None) as pool:
+        # Maximum works on the machine
+        pool.map(write_single_audio_to_hdf5, params)
+
+    print("Pack hdf5 time: {:.3f} s".format(time.time() - pack_hdf5s_time))
+
+
+def write_single_audio_to_hdf5(param: List) -> NoReturn:
+    r"""Write single audio into hdf5 file."""
+
+    (
+        audio_index,
+        audio_name,
+        speech_path,
+        mixture_path,
+        mono,
+        sample_rate,
+        hdf5_path,
+    ) = param
+
+    with h5py.File(hdf5_path, "w") as hf:
+
+        hf.attrs.create("audio_name", data=audio_name, dtype="S100")
+        hf.attrs.create("sample_rate", data=sample_rate, dtype=np.int32)
+
+        speech = load_audio(audio_path=speech_path, mono=mono, sample_rate=sample_rate)
+        # speech: (channels_num, audio_samples)
+
+        mixture = load_audio(
+            audio_path=mixture_path, mono=mono, sample_rate=sample_rate
+        )
+        # mixture: (channels_num, audio_samples)
+
+        noise = mixture - speech
+        # noise: (channels_num, audio_samples)
+
+        hf.create_dataset(name='speech', data=float32_to_int16(speech), dtype=np.int16)
+        hf.create_dataset(name='noise', data=float32_to_int16(noise), dtype=np.int16)
+
+    print('{} Write hdf5 to {}'.format(audio_index, hdf5_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory of the Voicebank-Demand dataset.",
+    )
+    parser.add_argument("--split", type=str, required=True, choices=["train", "test"])
+    parser.add_argument(
+        "--hdf5s_dir",
+        type=str,
+        required=True,
+        help="Directory to write out hdf5 files.",
+    )
+    parser.add_argument("--sample_rate", type=int, required=True, help="Sample rate.")
+    parser.add_argument(
+        "--channels", type=int, required=True, help="Use 1 for mono, 2 for stereo."
+    )
+
+    # Parse arguments.
+    args = parser.parse_args()
+
+    # Pack audios into hdf5 files.
+    pack_audios_to_hdf5s(args)

bytesep/inference.py

@@ -0,0 +1,404 @@
+import sys
+sys.path.append('.')
+import argparse
+import os
+import time
+from typing import Dict
+import pathlib
+
+import librosa
+import numpy as np
+import soundfile
+import torch
+import torch.nn as nn
+
+from bytesep.models.lightning_modules import get_model_class
+from bytesep.utils import read_yaml
+
+
+class Separator:
+    def __init__(
+        self, model: nn.Module, segment_samples: int, batch_size: int, device: str
+    ):
+        r"""Separate to separate an audio clip into a target source.
+
+        Args:
+            model: nn.Module, trained model
+            segment_samples: int, length of segments to be input to a model, e.g., 44100*30
+            batch_size, int, e.g., 12
+            device: str, e.g., 'cuda'
+        """
+        self.model = model
+        self.segment_samples = segment_samples
+        self.batch_size = batch_size
+        self.device = device
+
+    def separate(self, input_dict: Dict) -> np.array:
+        r"""Separate an audio clip into a target source.
+
+        Args:
+            input_dict: dict, e.g., {
+                waveform: (channels_num, audio_samples),
+                ...,
+            }
+
+        Returns:
+            sep_audio: (channels_num, audio_samples) | (target_sources_num, channels_num, audio_samples)
+        """
+        audio = input_dict['waveform']
+
+        audio_samples = audio.shape[-1]
+
+        # Pad the audio with zero in the end so that the length of audio can be
+        # evenly divided by segment_samples.
+        audio = self.pad_audio(audio)
+
+        # Enframe long audio into segments.
+        segments = self.enframe(audio, self.segment_samples)
+        # (segments_num, channels_num, segment_samples)
+
+        segments_input_dict = {'waveform': segments}
+
+        if 'condition' in input_dict.keys():
+            segments_num = len(segments)
+            segments_input_dict['condition'] = np.tile(
+                input_dict['condition'][None, :], (segments_num, 1)
+            )
+            # (batch_size, segments_num)
+
+        # Separate in mini-batches.
+        sep_segments = self._forward_in_mini_batches(
+            self.model, segments_input_dict, self.batch_size
+        )['waveform']
+        # (segments_num, channels_num, segment_samples)
+
+        # Deframe segments into long audio.
+        sep_audio = self.deframe(sep_segments)
+        # (channels_num, padded_audio_samples)
+
+        sep_audio = sep_audio[:, 0:audio_samples]
+        # (channels_num, audio_samples)
+
+        return sep_audio
+
+    def pad_audio(self, audio: np.array) -> np.array:
+        r"""Pad the audio with zero in the end so that the length of audio can
+        be evenly divided by segment_samples.
+
+        Args:
+            audio: (channels_num, audio_samples)
+
+        Returns:
+            padded_audio: (channels_num, audio_samples)
+        """
+        channels_num, audio_samples = audio.shape
+
+        # Number of segments
+        segments_num = int(np.ceil(audio_samples / self.segment_samples))
+
+        pad_samples = segments_num * self.segment_samples - audio_samples
+
+        padded_audio = np.concatenate(
+            (audio, np.zeros((channels_num, pad_samples))), axis=1
+        )
+        # (channels_num, padded_audio_samples)
+
+        return padded_audio
+
+    def enframe(self, audio: np.array, segment_samples: int) -> np.array:
+        r"""Enframe long audio into segments.
+
+        Args:
+            audio: (channels_num, audio_samples)
+            segment_samples: int
+
+        Returns:
+            segments: (segments_num, channels_num, segment_samples)
+        """
+        audio_samples = audio.shape[1]
+        assert audio_samples % segment_samples == 0
+
+        hop_samples = segment_samples // 2
+        segments = []
+
+        pointer = 0
+        while pointer + segment_samples <= audio_samples:
+            segments.append(audio[:, pointer : pointer + segment_samples])
+            pointer += hop_samples
+
+        segments = np.array(segments)
+
+        return segments
+
+    def deframe(self, segments: np.array) -> np.array:
+        r"""Deframe segments into long audio.
+
+        Args:
+            segments: (segments_num, channels_num, segment_samples)
+
+        Returns:
+            output: (channels_num, audio_samples)
+        """
+        (segments_num, _, segment_samples) = segments.shape
+
+        if segments_num == 1:
+            return segments[0]
+
+        assert self._is_integer(segment_samples * 0.25)
+        assert self._is_integer(segment_samples * 0.75)
+
+        output = []
+
+        output.append(segments[0, :, 0 : int(segment_samples * 0.75)])
+
+        for i in range(1, segments_num - 1):
+            output.append(
+                segments[
+                    i, :, int(segment_samples * 0.25) : int(segment_samples * 0.75)
+                ]
+            )
+
+        output.append(segments[-1, :, int(segment_samples * 0.25) :])
+
+        output = np.concatenate(output, axis=-1)
+
+        return output
+
+    def _is_integer(self, x: float) -> bool:
+        if x - int(x) < 1e-10:
+            return True
+        else:
+            return False
+
+    def _forward_in_mini_batches(
+        self, model: nn.Module, segments_input_dict: Dict, batch_size: int
+    ) -> Dict:
+        r"""Forward data to model in mini-batch.
+
+        Args:
+            model: nn.Module
+            segments_input_dict: dict, e.g., {
+                'waveform': (segments_num, channels_num, segment_samples),
+                ...,
+            }
+            batch_size: int
+
+        Returns:
+            output_dict: dict, e.g. {
+                'waveform': (segments_num, channels_num, segment_samples),
+            }
+        """
+        output_dict = {}
+
+        pointer = 0
+        segments_num = len(segments_input_dict['waveform'])
+
+        while True:
+            if pointer >= segments_num:
+                break
+
+            batch_input_dict = {}
+
+            for key in segments_input_dict.keys():
+                batch_input_dict[key] = torch.Tensor(
+                    segments_input_dict[key][pointer : pointer + batch_size]
+                ).to(self.device)
+
+            pointer += batch_size
+
+            with torch.no_grad():
+                model.eval()
+                batch_output_dict = model(batch_input_dict)
+
+            for key in batch_output_dict.keys():
+                self._append_to_dict(
+                    output_dict, key, batch_output_dict[key].data.cpu().numpy()
+                )
+
+        for key in output_dict.keys():
+            output_dict[key] = np.concatenate(output_dict[key], axis=0)
+
+        return output_dict
+
+    def _append_to_dict(self, dict, key, value):
+        if key in dict.keys():
+            dict[key].append(value)
+        else:
+            dict[key] = [value]
+
+
+class SeparatorWrapper:
+    def __init__(
+        self, source_type='vocals', model=None, checkpoint_path=None, device='cuda'
+    ):
+
+        input_channels = 2
+        target_sources_num = 1
+        model_type = "ResUNet143_Subbandtime"
+        segment_samples = 44100 * 10
+        batch_size = 1
+
+        self.checkpoint_path = self.download_checkpoints(checkpoint_path, source_type)
+
+        if device == 'cuda' and torch.cuda.is_available():
+            self.device = 'cuda'
+        else:
+            self.device = 'cpu'
+
+        # Get model class.
+        Model = get_model_class(model_type)
+
+        # Create model.
+        self.model = Model(
+            input_channels=input_channels, target_sources_num=target_sources_num
+        )
+
+        # Load checkpoint.
+        checkpoint = torch.load(self.checkpoint_path, map_location='cpu')
+        self.model.load_state_dict(checkpoint["model"])
+
+        # Move model to device.
+        self.model.to(self.device)
+
+        # Create separator.
+        self.separator = Separator(
+            model=self.model,
+            segment_samples=segment_samples,
+            batch_size=batch_size,
+            device=self.device,
+        )
+
+    def download_checkpoints(self, checkpoint_path, source_type):
+
+        if source_type == "vocals":
+            checkpoint_bare_name = "resunet143_subbtandtime_vocals_8.8dB_350k_steps"
+
+        elif source_type == "accompaniment":
+            checkpoint_bare_name = (
+                "resunet143_subbtandtime_accompaniment_16.4dB_350k_steps.pth"
+            )
+
+        else:
+            raise NotImplementedError
+
+        if not checkpoint_path:
+            checkpoint_path = '{}/bytesep_data/{}.pth'.format(
+                str(pathlib.Path.home()), checkpoint_bare_name
+            )
+
+        print('Checkpoint path: {}'.format(checkpoint_path))
+
+        if (
+            not os.path.exists(checkpoint_path)
+            or os.path.getsize(checkpoint_path) < 4e8
+        ):
+
+            os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
+
+            zenodo_dir = "https://zenodo.org/record/5507029/files"
+            zenodo_path = os.path.join(
+                zenodo_dir, "{}?download=1".format(checkpoint_bare_name)
+            )
+
+            os.system('wget -O "{}" "{}"'.format(checkpoint_path, zenodo_path))
+
+        return checkpoint_path
+
+    def separate(self, audio):
+
+        input_dict = {'waveform': audio}
+
+        sep_wav = self.separator.separate(input_dict)
+
+        return sep_wav
+
+
+def inference(args):
+
+    # Need to use torch.distributed if models contain inplace_abn.abn.InPlaceABNSync.
+    import torch.distributed as dist
+
+    dist.init_process_group(
+        'gloo', init_method='file:///tmp/somefile', rank=0, world_size=1
+    )
+
+    # Arguments & parameters
+    config_yaml = args.config_yaml
+    checkpoint_path = args.checkpoint_path
+    audio_path = args.audio_path
+    output_path = args.output_path
+    device = (
+        torch.device('cuda')
+        if args.cuda and torch.cuda.is_available()
+        else torch.device('cpu')
+    )
+
+    configs = read_yaml(config_yaml)
+    sample_rate = configs['train']['sample_rate']
+    input_channels = configs['train']['channels']
+    target_source_types = configs['train']['target_source_types']
+    target_sources_num = len(target_source_types)
+    model_type = configs['train']['model_type']
+
+    segment_samples = int(30 * sample_rate)
+    batch_size = 1
+
+    print("Using {} for separating ..".format(device))
+
+    # paths
+    if os.path.dirname(output_path) != "":
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # Get model class.
+    Model = get_model_class(model_type)
+
+    # Create model.
+    model = Model(input_channels=input_channels, target_sources_num=target_sources_num)
+
+    # Load checkpoint.
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    model.load_state_dict(checkpoint["model"])
+
+    # Move model to device.
+    model.to(device)
+
+    # Create separator.
+    separator = Separator(
+        model=model,
+        segment_samples=segment_samples,
+        batch_size=batch_size,
+        device=device,
+    )
+
+    # Load audio.
+    audio, _ = librosa.load(audio_path, sr=sample_rate, mono=False)
+
+    # audio = audio[None, :]
+
+    input_dict = {'waveform': audio}
+
+    # Separate
+    separate_time = time.time()
+
+    sep_wav = separator.separate(input_dict)
+    # (channels_num, audio_samples)
+
+    print('Separate time: {:.3f} s'.format(time.time() - separate_time))
+
+    # Write out separated audio.
+    soundfile.write(file='_zz.wav', data=sep_wav.T, samplerate=sample_rate)
+    os.system("ffmpeg -y -loglevel panic -i _zz.wav {}".format(output_path))
+    print('Write out to {}'.format(output_path))
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument("--config_yaml", type=str, required=True)
+    parser.add_argument("--checkpoint_path", type=str, required=True)
+    parser.add_argument("--audio_path", type=str, required=True)
+    parser.add_argument("--output_path", type=str, required=True)
+    parser.add_argument("--cuda", action='store_true', default=True)
+
+    args = parser.parse_args()
+    inference(args)

bytesep/inference_many.py

@@ -0,0 +1,163 @@
+import argparse
+import os
+import pathlib
+import time
+from typing import NoReturn
+
+import librosa
+import numpy as np
+import soundfile
+import torch
+
+from bytesep.inference import Separator
+from bytesep.models.lightning_modules import get_model_class
+from bytesep.utils import read_yaml
+
+
+def inference(args) -> NoReturn:
+    r"""Separate all audios in a directory.
+
+    Args:
+        config_yaml: str, the config file of a model being trained
+        checkpoint_path: str, the path of checkpoint to be loaded
+        audios_dir: str, the directory of audios to be separated
+        output_dir: str, the directory to write out separated audios
+        scale_volume: bool, if True then the volume is scaled to the maximum value of 1.
+
+    Returns:
+        NoReturn
+    """
+
+    # Arguments & parameters
+    config_yaml = args.config_yaml
+    checkpoint_path = args.checkpoint_path
+    audios_dir = args.audios_dir
+    output_dir = args.output_dir
+    scale_volume = args.scale_volume
+    device = (
+        torch.device('cuda')
+        if args.cuda and torch.cuda.is_available()
+        else torch.device('cpu')
+    )
+
+    configs = read_yaml(config_yaml)
+    sample_rate = configs['train']['sample_rate']
+    input_channels = configs['train']['channels']
+    target_source_types = configs['train']['target_source_types']
+    target_sources_num = len(target_source_types)
+    model_type = configs['train']['model_type']
+    mono = input_channels == 1
+
+    segment_samples = int(30 * sample_rate)
+    batch_size = 1
+    device = "cuda"
+
+    models_contains_inplaceabn = True
+
+    # Need to use torch.distributed if models contain inplace_abn.abn.InPlaceABNSync.
+    if models_contains_inplaceabn:
+
+        import torch.distributed as dist
+
+        dist.init_process_group(
+            'gloo', init_method='file:///tmp/somefile', rank=0, world_size=1
+        )
+
+    print("Using {} for separating ..".format(device))
+
+    # paths
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Get model class.
+    Model = get_model_class(model_type)
+
+    # Create model.
+    model = Model(input_channels=input_channels, target_sources_num=target_sources_num)
+
+    # Load checkpoint.
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+    model.load_state_dict(checkpoint["model"])
+
+    # Move model to device.
+    model.to(device)
+
+    # Create separator.
+    separator = Separator(
+        model=model,
+        segment_samples=segment_samples,
+        batch_size=batch_size,
+        device=device,
+    )
+
+    audio_names = sorted(os.listdir(audios_dir))
+
+    for audio_name in audio_names:
+        audio_path = os.path.join(audios_dir, audio_name)
+
+        # Load audio.
+        audio, _ = librosa.load(audio_path, sr=sample_rate, mono=mono)
+
+        if audio.ndim == 1:
+            audio = audio[None, :]
+
+        input_dict = {'waveform': audio}
+
+        # Separate
+        separate_time = time.time()
+
+        sep_wav = separator.separate(input_dict)
+        # (channels_num, audio_samples)
+
+        print('Separate time: {:.3f} s'.format(time.time() - separate_time))
+
+        # Write out separated audio.
+        if scale_volume:
+            sep_wav /= np.max(np.abs(sep_wav))
+
+        soundfile.write(file='_zz.wav', data=sep_wav.T, samplerate=sample_rate)
+
+        output_path = os.path.join(
+            output_dir, '{}.mp3'.format(pathlib.Path(audio_name).stem)
+        )
+        os.system('ffmpeg -y -loglevel panic -i _zz.wav "{}"'.format(output_path))
+        print('Write out to {}'.format(output_path))
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--config_yaml",
+        type=str,
+        required=True,
+        help="The config file of a model being trained.",
+    )
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        required=True,
+        help="The path of checkpoint to be loaded.",
+    )
+    parser.add_argument(
+        "--audios_dir",
+        type=str,
+        required=True,
+        help="The directory of audios to be separated.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="The directory to write out separated audios.",
+    )
+    parser.add_argument(
+        '--scale_volume',
+        action='store_true',
+        default=False,
+        help="set to True if separated audios are scaled to the maximum value of 1.",
+    )
+    parser.add_argument("--cuda", action='store_true', default=True)
+
+    args = parser.parse_args()
+
+    inference(args)

bytesep/losses.py

@@ -0,0 +1,106 @@
+import math
+from typing import Callable
+
+import torch
+import torch.nn as nn
+from torchlibrosa.stft import STFT
+
+from bytesep.models.pytorch_modules import Base
+
+
+def l1(output: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor:
+    r"""L1 loss.
+
+    Args:
+        output: torch.Tensor
+        target: torch.Tensor
+
+    Returns:
+        loss: torch.float
+    """
+    return torch.mean(torch.abs(output - target))
+
+
+def l1_wav(output: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor:
+    r"""L1 loss in the time-domain.
+
+    Args:
+        output: torch.Tensor
+        target: torch.Tensor
+
+    Returns:
+        loss: torch.float
+    """
+    return l1(output, target)
+
+
+class L1_Wav_L1_Sp(nn.Module, Base):
+    def __init__(self):
+        r"""L1 loss in the time-domain and L1 loss on the spectrogram."""
+        super(L1_Wav_L1_Sp, self).__init__()
+
+        self.window_size = 2048
+        hop_size = 441
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+
+        self.stft = STFT(
+            n_fft=self.window_size,
+            hop_length=hop_size,
+            win_length=self.window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+    def __call__(
+        self, output: torch.Tensor, target: torch.Tensor, **kwargs
+    ) -> torch.Tensor:
+        r"""L1 loss in the time-domain and on the spectrogram.
+
+        Args:
+            output: torch.Tensor
+            target: torch.Tensor
+
+        Returns:
+            loss: torch.float
+        """
+
+        # L1 loss in the time-domain.
+        wav_loss = l1_wav(output, target)
+
+        # L1 loss on the spectrogram.
+        sp_loss = l1(
+            self.wav_to_spectrogram(output, eps=1e-8),
+            self.wav_to_spectrogram(target, eps=1e-8),
+        )
+
+        # sp_loss /= math.sqrt(self.window_size)
+        # sp_loss *= 1.
+
+        # Total loss.
+        return wav_loss + sp_loss
+
+        return sp_loss
+
+
+def get_loss_function(loss_type: str) -> Callable:
+    r"""Get loss function.
+
+    Args:
+        loss_type: str
+
+    Returns:
+        loss function: Callable
+    """
+
+    if loss_type == "l1_wav":
+        return l1_wav
+
+    elif loss_type == "l1_wav_l1_sp":
+        return L1_Wav_L1_Sp()
+
+    else:
+        raise NotImplementedError

bytesep/models/conditional_unet.py

@@ -0,0 +1,496 @@
+import math
+from typing import List
+
+import numpy as np
+import matplotlib.pyplot as plt
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.optim.lr_scheduler import LambdaLR
+from torchlibrosa.stft import STFT, ISTFT, magphase
+
+from bytesep.models.pytorch_modules import (
+    Base,
+    init_bn,
+    init_embedding,
+    init_layer,
+    act,
+    Subband,
+)
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        condition_size,
+        kernel_size,
+        activation,
+        momentum,
+    ):
+        super(ConvBlock, self).__init__()
+
+        self.activation = activation
+        padding = (kernel_size[0] // 2, kernel_size[1] // 2)
+
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.beta1 = nn.Linear(condition_size, out_channels, bias=True)
+        self.beta2 = nn.Linear(condition_size, out_channels, bias=True)
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+        init_embedding(self.beta1)
+        init_embedding(self.beta2)
+
+    def forward(self, x, condition):
+
+        b1 = self.beta1(condition)[:, :, None, None]
+        b2 = self.beta2(condition)[:, :, None, None]
+
+        x = act(self.bn1(self.conv1(x)) + b1, self.activation)
+        x = act(self.bn2(self.conv2(x)) + b2, self.activation)
+        return x
+
+
+class EncoderBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        condition_size,
+        kernel_size,
+        downsample,
+        activation,
+        momentum,
+    ):
+        super(EncoderBlock, self).__init__()
+
+        self.conv_block = ConvBlock(
+            in_channels, out_channels, condition_size, kernel_size, activation, momentum
+        )
+        self.downsample = downsample
+
+    def forward(self, x, condition):
+        encoder = self.conv_block(x, condition)
+        encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample)
+        return encoder_pool, encoder
+
+
+class DecoderBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        condition_size,
+        kernel_size,
+        upsample,
+        activation,
+        momentum,
+    ):
+        super(DecoderBlock, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = upsample
+        self.activation = activation
+
+        self.conv1 = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=self.stride,
+            stride=self.stride,
+            padding=(0, 0),
+            bias=False,
+            dilation=(1, 1),
+        )
+
+        self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv_block2 = ConvBlock(
+            out_channels * 2,
+            out_channels,
+            condition_size,
+            kernel_size,
+            activation,
+            momentum,
+        )
+
+        self.beta1 = nn.Linear(condition_size, out_channels, bias=True)
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_layer(self.conv1)
+        init_bn(self.bn1)
+        init_embedding(self.beta1)
+
+    def forward(self, input_tensor, concat_tensor, condition):
+        b1 = self.beta1(condition)[:, :, None, None]
+        x = act(self.bn1(self.conv1(input_tensor)) + b1, self.activation)
+        x = torch.cat((x, concat_tensor), dim=1)
+        x = self.conv_block2(x, condition)
+        return x
+
+
+class ConditionalUNet(nn.Module, Base):
+    def __init__(self, input_channels, target_sources_num):
+        super(ConditionalUNet, self).__init__()
+
+        self.input_channels = input_channels
+        condition_size = target_sources_num
+        self.output_sources_num = 1
+
+        window_size = 2048
+        hop_size = 441
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+        activation = "relu"
+        momentum = 0.01
+
+        self.subbands_num = 4
+        self.K = 3  # outputs: |M|, cos∠M, sin∠M
+
+        self.downsample_ratio = 2 ** 6  # This number equals 2^{#encoder_blcoks}
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.subband = Subband(subbands_num=self.subbands_num)
+
+        self.encoder_block1 = EncoderBlock(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlock(
+            in_channels=32,
+            out_channels=64,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlock(
+            in_channels=64,
+            out_channels=128,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlock(
+            in_channels=128,
+            out_channels=256,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlock(
+            in_channels=256,
+            out_channels=384,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlock(
+            in_channels=384,
+            out_channels=384,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7 = ConvBlock(
+            in_channels=384,
+            out_channels=384,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlock(
+            in_channels=384,
+            out_channels=256,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlock(
+            in_channels=256,
+            out_channels=128,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlock(
+            in_channels=128,
+            out_channels=64,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block6 = DecoderBlock(
+            in_channels=64,
+            out_channels=32,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = ConvBlock(
+            in_channels=32,
+            out_channels=32,
+            condition_size=condition_size,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=input_channels
+            * self.subbands_num
+            * self.output_sources_num
+            * self.K,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(self, x, sp, sin_in, cos_in, audio_length):
+
+        batch_size, _, time_steps, freq_bins = x.shape
+
+        x = x.reshape(
+            batch_size,
+            self.output_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, output_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        # mask_cos, mask_sin: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos, out_sin: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag)
+        # out_mag: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.output_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        wav_out = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * output_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        wav_out = wav_out.reshape(
+            batch_size, self.output_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, output_sources_num * input_channels, segments_num)
+
+        return wav_out
+
+    def forward(self, input_dict):
+        """
+        Args:
+          input: (batch_size, segment_samples, channels_num)
+
+        Outputs:
+          output_dict: {
+            'wav': (batch_size, segment_samples, channels_num),
+            'sp': (batch_size, channels_num, time_steps, freq_bins)}
+        """
+
+        mixture = input_dict['waveform']
+        condition = input_dict['condition']
+
+        sp, cos_in, sin_in = self.wav_to_spectrogram_phase(mixture)
+        """(batch_size, channels_num, time_steps, freq_bins)"""
+
+        # Batch normalization
+        x = sp.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        """(batch_size, chanenls, time_steps, freq_bins)"""
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        """(batch_size, channels, padded_time_steps, freq_bins)"""
+
+        # Let frequency bins be evenly divided by 2, e.g., 513 -> 512
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, channels, T, F)
+
+        x = self.subband.analysis(x)
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(
+            x, condition
+        )  # x1_pool: (bs, 32, T / 2, F / 2)
+        (x2_pool, x2) = self.encoder_block2(
+            x1_pool, condition
+        )  # x2_pool: (bs, 64, T / 4, F / 4)
+        (x3_pool, x3) = self.encoder_block3(
+            x2_pool, condition
+        )  # x3_pool: (bs, 128, T / 8, F / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool, condition
+        )  # x4_pool: (bs, 256, T / 16, F / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool, condition
+        )  # x5_pool: (bs, 512, T / 32, F / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool, condition
+        )  # x6_pool: (bs, 1024, T / 64, F / 64)
+        x_center = self.conv_block7(x6_pool, condition)  # (bs, 2048, T / 64, F / 64)
+        x7 = self.decoder_block1(x_center, x6, condition)  # (bs, 1024, T / 32, F / 32)
+        x8 = self.decoder_block2(x7, x5, condition)  # (bs, 512, T / 16, F / 16)
+        x9 = self.decoder_block3(x8, x4, condition)  # (bs, 256, T / 8, F / 8)
+        x10 = self.decoder_block4(x9, x3, condition)  # (bs, 128, T / 4, F / 4)
+        x11 = self.decoder_block5(x10, x2, condition)  # (bs, 64, T / 2, F / 2)
+        x12 = self.decoder_block6(x11, x1, condition)  # (bs, 32, T, F)
+        x = self.after_conv_block1(x12, condition)  # (bs, 32, T, F)
+        x = self.after_conv2(x)
+        # (batch_size, input_channles * subbands_num * targets_num * k, T, F // subbands_num)
+
+        x = self.subband.synthesis(x)
+        # (batch_size, input_channles * targets_num * K, T, F)
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 1024 -> 1025.
+        x = x[:, :, 0:origin_len, :]  # (bs, feature_maps, T, F)
+
+        audio_length = mixture.shape[2]
+
+        separated_audio = self.feature_maps_to_wav(x, sp, sin_in, cos_in, audio_length)
+        # separated_audio: (batch_size, output_sources_num * input_channels, segments_num)
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/models/lightning_modules.py

@@ -0,0 +1,188 @@
+from typing import Any, Callable, Dict
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import LambdaLR
+
+
+class LitSourceSeparation(pl.LightningModule):
+    def __init__(
+        self,
+        batch_data_preprocessor,
+        model: nn.Module,
+        loss_function: Callable,
+        optimizer_type: str,
+        learning_rate: float,
+        lr_lambda: Callable,
+    ):
+        r"""Pytorch Lightning wrapper of PyTorch model, including forward,
+        optimization of model, etc.
+
+        Args:
+            batch_data_preprocessor: object, used for preparing inputs and
+                targets for training. E.g., BasicBatchDataPreprocessor is used
+                for preparing data in dictionary into tensor.
+            model: nn.Module
+            loss_function: function
+            learning_rate: float
+            lr_lambda: function
+        """
+        super().__init__()
+
+        self.batch_data_preprocessor = batch_data_preprocessor
+        self.model = model
+        self.optimizer_type = optimizer_type
+        self.loss_function = loss_function
+        self.learning_rate = learning_rate
+        self.lr_lambda = lr_lambda
+
+    def training_step(self, batch_data_dict: Dict, batch_idx: int) -> torch.float:
+        r"""Forward a mini-batch data to model, calculate loss function, and
+        train for one step. A mini-batch data is evenly distributed to multiple
+        devices (if there are) for parallel training.
+
+        Args:
+            batch_data_dict: e.g. {
+                'vocals': (batch_size, channels_num, segment_samples),
+                'accompaniment': (batch_size, channels_num, segment_samples),
+                'mixture': (batch_size, channels_num, segment_samples)
+            }
+            batch_idx: int
+
+        Returns:
+            loss: float, loss function of this mini-batch
+        """
+        input_dict, target_dict = self.batch_data_preprocessor(batch_data_dict)
+        # input_dict: {
+        #     'waveform': (batch_size, channels_num, segment_samples),
+        #     (if_exist) 'condition': (batch_size, channels_num),
+        # }
+        # target_dict: {
+        #     'waveform': (batch_size, target_sources_num * channels_num, segment_samples),
+        # }
+
+        # Forward.
+        self.model.train()
+
+        output_dict = self.model(input_dict)
+        # output_dict: {
+        #     'waveform': (batch_size, target_sources_num * channels_num, segment_samples),
+        # }
+
+        outputs = output_dict['waveform']
+        # outputs:, e.g, (batch_size, target_sources_num * channels_num, segment_samples)
+
+        # Calculate loss.
+        loss = self.loss_function(
+            output=outputs,
+            target=target_dict['waveform'],
+            mixture=input_dict['waveform'],
+        )
+
+        return loss
+
+    def configure_optimizers(self) -> Any:
+        r"""Configure optimizer."""
+
+        if self.optimizer_type == "Adam":
+            optimizer = optim.Adam(
+                self.model.parameters(),
+                lr=self.learning_rate,
+                betas=(0.9, 0.999),
+                eps=1e-08,
+                weight_decay=0.0,
+                amsgrad=True,
+            )
+
+        elif self.optimizer_type == "AdamW":
+            optimizer = optim.AdamW(
+                self.model.parameters(),
+                lr=self.learning_rate,
+                betas=(0.9, 0.999),
+                eps=1e-08,
+                weight_decay=0.0,
+                amsgrad=True,
+            )
+
+        else:
+            raise NotImplementedError
+
+        scheduler = {
+            'scheduler': LambdaLR(optimizer, self.lr_lambda),
+            'interval': 'step',
+            'frequency': 1,
+        }
+
+        return [optimizer], [scheduler]
+
+
+def get_model_class(model_type):
+    r"""Get model.
+
+    Args:
+        model_type: str, e.g., 'ResUNet143_DecouplePlusInplaceABN'
+
+    Returns:
+        nn.Module
+    """
+    if model_type == 'ResUNet143_DecouplePlusInplaceABN_ISMIR2021':
+        from bytesep.models.resunet_ismir2021 import (
+            ResUNet143_DecouplePlusInplaceABN_ISMIR2021,
+        )
+
+        return ResUNet143_DecouplePlusInplaceABN_ISMIR2021
+
+    elif model_type == 'UNet':
+        from bytesep.models.unet import UNet
+
+        return UNet
+
+    elif model_type == 'UNetSubbandTime':
+        from bytesep.models.unet_subbandtime import UNetSubbandTime
+
+        return UNetSubbandTime
+
+    elif model_type == 'ResUNet143_Subbandtime':
+        from bytesep.models.resunet_subbandtime import ResUNet143_Subbandtime
+
+        return ResUNet143_Subbandtime
+
+    elif model_type == 'ResUNet143_DecouplePlus':
+        from bytesep.models.resunet import ResUNet143_DecouplePlus
+
+        return ResUNet143_DecouplePlus
+
+    elif model_type == 'ConditionalUNet':
+        from bytesep.models.conditional_unet import ConditionalUNet
+
+        return ConditionalUNet
+
+    elif model_type == 'LevelRNN':
+        from bytesep.models.levelrnn import LevelRNN
+
+        return LevelRNN
+
+    elif model_type == 'WavUNet':
+        from bytesep.models.wavunet import WavUNet
+
+        return WavUNet
+
+    elif model_type == 'WavUNetLevelRNN':
+        from bytesep.models.wavunet_levelrnn import WavUNetLevelRNN
+
+        return WavUNetLevelRNN
+
+    elif model_type == 'TTnet':
+        from bytesep.models.ttnet import TTnet
+
+        return TTnet
+
+    elif model_type == 'TTnetNoTransformer':
+        from bytesep.models.ttnet_no_transformer import TTnetNoTransformer
+
+        return TTnetNoTransformer
+
+    else:
+        raise NotImplementedError

bytesep/models/pytorch_modules.py

@@ -0,0 +1,204 @@
+from typing import List, NoReturn
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def init_embedding(layer: nn.Module) -> NoReturn:
+    r"""Initialize a Linear or Convolutional layer."""
+    nn.init.uniform_(layer.weight, -1.0, 1.0)
+
+    if hasattr(layer, 'bias'):
+        if layer.bias is not None:
+            layer.bias.data.fill_(0.0)
+
+
+def init_layer(layer: nn.Module) -> NoReturn:
+    r"""Initialize a Linear or Convolutional layer."""
+    nn.init.xavier_uniform_(layer.weight)
+
+    if hasattr(layer, "bias"):
+        if layer.bias is not None:
+            layer.bias.data.fill_(0.0)
+
+
+def init_bn(bn: nn.Module) -> NoReturn:
+    r"""Initialize a Batchnorm layer."""
+    bn.bias.data.fill_(0.0)
+    bn.weight.data.fill_(1.0)
+    bn.running_mean.data.fill_(0.0)
+    bn.running_var.data.fill_(1.0)
+
+
+def act(x: torch.Tensor, activation: str) -> torch.Tensor:
+
+    if activation == "relu":
+        return F.relu_(x)
+
+    elif activation == "leaky_relu":
+        return F.leaky_relu_(x, negative_slope=0.01)
+
+    elif activation == "swish":
+        return x * torch.sigmoid(x)
+
+    else:
+        raise Exception("Incorrect activation!")
+
+
+class Base:
+    def __init__(self):
+        r"""Base function for extracting spectrogram, cos, and sin, etc."""
+        pass
+
+    def spectrogram(self, input: torch.Tensor, eps: float = 0.0) -> torch.Tensor:
+        r"""Calculate spectrogram.
+
+        Args:
+            input: (batch_size, segments_num)
+            eps: float
+
+        Returns:
+            spectrogram: (batch_size, time_steps, freq_bins)
+        """
+        (real, imag) = self.stft(input)
+        return torch.clamp(real ** 2 + imag ** 2, eps, np.inf) ** 0.5
+
+    def spectrogram_phase(
+        self, input: torch.Tensor, eps: float = 0.0
+    ) -> List[torch.Tensor]:
+        r"""Calculate the magnitude, cos, and sin of the STFT of input.
+
+        Args:
+            input: (batch_size, segments_num)
+            eps: float
+
+        Returns:
+            mag: (batch_size, time_steps, freq_bins)
+            cos: (batch_size, time_steps, freq_bins)
+            sin: (batch_size, time_steps, freq_bins)
+        """
+        (real, imag) = self.stft(input)
+        mag = torch.clamp(real ** 2 + imag ** 2, eps, np.inf) ** 0.5
+        cos = real / mag
+        sin = imag / mag
+        return mag, cos, sin
+
+    def wav_to_spectrogram_phase(
+        self, input: torch.Tensor, eps: float = 1e-10
+    ) -> List[torch.Tensor]:
+        r"""Convert waveforms to magnitude, cos, and sin of STFT.
+
+        Args:
+            input: (batch_size, channels_num, segment_samples)
+            eps: float
+
+        Outputs:
+            mag: (batch_size, channels_num, time_steps, freq_bins)
+            cos: (batch_size, channels_num, time_steps, freq_bins)
+            sin: (batch_size, channels_num, time_steps, freq_bins)
+        """
+        batch_size, channels_num, segment_samples = input.shape
+
+        # Reshape input with shapes of (n, segments_num) to meet the
+        # requirements of the stft function.
+        x = input.reshape(batch_size * channels_num, segment_samples)
+
+        mag, cos, sin = self.spectrogram_phase(x, eps=eps)
+        # mag, cos, sin: (batch_size * channels_num, 1, time_steps, freq_bins)
+
+        _, _, time_steps, freq_bins = mag.shape
+        mag = mag.reshape(batch_size, channels_num, time_steps, freq_bins)
+        cos = cos.reshape(batch_size, channels_num, time_steps, freq_bins)
+        sin = sin.reshape(batch_size, channels_num, time_steps, freq_bins)
+
+        return mag, cos, sin
+
+    def wav_to_spectrogram(
+        self, input: torch.Tensor, eps: float = 1e-10
+    ) -> List[torch.Tensor]:
+
+        mag, cos, sin = self.wav_to_spectrogram_phase(input, eps)
+        return mag
+
+
+class Subband:
+    def __init__(self, subbands_num: int):
+        r"""Warning!! This class is not used!!
+
+        This class does not work as good as [1] which split subbands in the
+        time-domain. Please refere to [1] for formal implementation.
+
+        [1] Liu, Haohe, et al. "Channel-wise subband input for better voice and
+        accompaniment separation on high resolution music." arXiv preprint arXiv:2008.05216 (2020).
+
+        Args:
+            subbands_num: int, e.g., 4
+        """
+        self.subbands_num = subbands_num
+
+    def analysis(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Analysis time-frequency representation into subbands. Stack the
+        subbands along the channel axis.
+
+        Args:
+            x: (batch_size, channels_num, time_steps, freq_bins)
+
+        Returns:
+            output: (batch_size, channels_num * subbands_num, time_steps, freq_bins // subbands_num)
+        """
+        batch_size, channels_num, time_steps, freq_bins = x.shape
+
+        x = x.reshape(
+            batch_size,
+            channels_num,
+            time_steps,
+            self.subbands_num,
+            freq_bins // self.subbands_num,
+        )
+        # x: (batch_size, channels_num, time_steps, subbands_num, freq_bins // subbands_num)
+
+        x = x.transpose(2, 3)
+
+        output = x.reshape(
+            batch_size,
+            channels_num * self.subbands_num,
+            time_steps,
+            freq_bins // self.subbands_num,
+        )
+        # output: (batch_size, channels_num * subbands_num, time_steps, freq_bins // subbands_num)
+
+        return output
+
+    def synthesis(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Synthesis subband time-frequency representations into original
+        time-frequency representation.
+
+        Args:
+            x: (batch_size, channels_num * subbands_num, time_steps, freq_bins // subbands_num)
+
+        Returns:
+            output: (batch_size, channels_num, time_steps, freq_bins)
+        """
+        batch_size, subband_channels_num, time_steps, subband_freq_bins = x.shape
+
+        channels_num = subband_channels_num // self.subbands_num
+        freq_bins = subband_freq_bins * self.subbands_num
+
+        x = x.reshape(
+            batch_size,
+            channels_num,
+            self.subbands_num,
+            time_steps,
+            subband_freq_bins,
+        )
+        # x: (batch_size, channels_num, subbands_num, time_steps, freq_bins // subbands_num)
+
+        x = x.transpose(2, 3)
+        # x: (batch_size, channels_num, time_steps, subbands_num, freq_bins // subbands_num)
+
+        output = x.reshape(batch_size, channels_num, time_steps, freq_bins)
+        # x: (batch_size, channels_num, time_steps, freq_bins)
+
+        return output

bytesep/models/resunet.py

@@ -0,0 +1,516 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchlibrosa.stft import ISTFT, STFT, magphase
+
+from bytesep.models.pytorch_modules import Base, Subband, act, init_bn, init_layer
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, activation, momentum):
+        r"""Residual block."""
+        super(ConvBlockRes, self).__init__()
+
+        self.activation = activation
+        padding = [kernel_size[0] // 2, kernel_size[1] // 2]
+
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(1, 1),
+                stride=(1, 1),
+                padding=(0, 0),
+            )
+
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+
+        if self.is_shortcut:
+            init_layer(self.shortcut)
+
+    def forward(self, x):
+        origin = x
+        x = self.conv1(act(self.bn1(x), self.activation))
+        x = self.conv2(act(self.bn2(x), self.activation))
+
+        if self.is_shortcut:
+            return self.shortcut(origin) + x
+        else:
+            return origin + x
+
+
+class EncoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, downsample, activation, momentum
+    ):
+        r"""Encoder block, contains 8 convolutional layers."""
+        super(EncoderBlockRes4B, self).__init__()
+
+        self.conv_block1 = ConvBlockRes(
+            in_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block2 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.downsample = downsample
+
+    def forward(self, x):
+        encoder = self.conv_block1(x)
+        encoder = self.conv_block2(encoder)
+        encoder = self.conv_block3(encoder)
+        encoder = self.conv_block4(encoder)
+        encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample)
+        return encoder_pool, encoder
+
+
+class DecoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, upsample, activation, momentum
+    ):
+        r"""Decoder block, contains 1 transpose convolutional and 8 convolutional layers."""
+        super(DecoderBlockRes4B, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = upsample
+        self.activation = activation
+
+        self.conv1 = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=self.stride,
+            stride=self.stride,
+            padding=(0, 0),
+            bias=False,
+            dilation=(1, 1),
+        )
+
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.conv_block2 = ConvBlockRes(
+            out_channels * 2, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block5 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_layer(self.conv1)
+
+    def forward(self, input_tensor, concat_tensor):
+        x = self.conv1(act(self.bn1(input_tensor), self.activation))
+        x = torch.cat((x, concat_tensor), dim=1)
+        x = self.conv_block2(x)
+        x = self.conv_block3(x)
+        x = self.conv_block4(x)
+        x = self.conv_block5(x)
+        return x
+
+
+class ResUNet143_DecouplePlus(nn.Module, Base):
+    def __init__(self, input_channels, target_sources_num):
+        super(ResUNet143_DecouplePlus, self).__init__()
+
+        self.input_channels = input_channels
+        self.target_sources_num = target_sources_num
+
+        window_size = 2048
+        hop_size = 441
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+        activation = "relu"
+        momentum = 0.01
+
+        self.subbands_num = 4
+        self.K = 4  # outputs: |M|, cos∠M, sin∠M, |M2|
+
+        self.downsample_ratio = 2 ** 6  # This number equals 2^{#encoder_blcoks}
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.subband = Subband(subbands_num=self.subbands_num)
+
+        self.encoder_block1 = EncoderBlockRes4B(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlockRes4B(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlockRes4B(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlockRes4B(
+            in_channels=256,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7a = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7b = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7c = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7d = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=256,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlockRes4B(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlockRes4B(
+            in_channels=128,
+            out_channels=64,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block6 = DecoderBlockRes4B(
+            in_channels=64,
+            out_channels=32,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=input_channels
+            * self.subbands_num
+            * target_sources_num
+            * self.K,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(
+        self,
+        input_tensor: torch.Tensor,
+        sp: torch.Tensor,
+        sin_in: torch.Tensor,
+        cos_in: torch.Tensor,
+        audio_length: int,
+    ) -> torch.Tensor:
+        r"""Convert feature maps to waveform.
+
+        Args:
+            input_tensor: (batch_size, feature_maps, time_steps, freq_bins)
+            sp: (batch_size, feature_maps, time_steps, freq_bins)
+            sin_in: (batch_size, feature_maps, time_steps, freq_bins)
+            cos_in: (batch_size, feature_maps, time_steps, freq_bins)
+
+        Outputs:
+            waveform: (batch_size, target_sources_num * input_channels, segment_samples)
+        """
+        batch_size, _, time_steps, freq_bins = input_tensor.shape
+
+        x = input_tensor.reshape(
+            batch_size,
+            self.target_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, target_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        linear_mag = x[:, :, :, 3, :, :]
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        # mask_cos, mask_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+        # out_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag + linear_mag)
+        # out_mag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.target_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        x = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * target_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        waveform = x.reshape(
+            batch_size, self.target_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, target_sources_num * input_channels, segments_num)
+
+        return waveform
+
+    def forward(self, input_dict):
+        r"""
+        Args:
+            input: (batch_size, channels_num, segment_samples)
+
+        Outputs:
+            output_dict: {
+                'wav': (batch_size, channels_num, segment_samples)
+            }
+        """
+        mixtures = input_dict['waveform']
+        # (batch_size, input_channels, segment_samples)
+
+        mag, cos_in, sin_in = self.wav_to_spectrogram_phase(mixtures)
+        # mag, cos_in, sin_in: (batch_size, input_channels, time_steps, freq_bins)
+
+        # Batch normalize on individual frequency bins.
+        x = mag.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        """(batch_size, input_channels, time_steps, freq_bins)"""
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        """(batch_size, input_channels, padded_time_steps, freq_bins)"""
+
+        # Let frequency bins be evenly divided by 2, e.g., 1025 -> 1024
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, input_channels, T, F)
+
+        x = self.subband.analysis(x)
+        # (bs, input_channels, T, F'), where F' = F // subbands_num
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F / 2)
+        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F / 4)
+        (x3_pool, x3) = self.encoder_block3(x2_pool)  # x3_pool: (bs, 128, T / 8, F / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool
+        )  # x4_pool: (bs, 256, T / 16, F / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool
+        )  # x5_pool: (bs, 384, T / 32, F / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool
+        )  # x6_pool: (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7a(x6_pool)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7b(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7c(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7d(x_center)  # (bs, 384, T / 32, F / 64)
+        x7 = self.decoder_block1(x_center, x6)  # (bs, 384, T / 32, F / 32)
+        x8 = self.decoder_block2(x7, x5)  # (bs, 384, T / 16, F / 16)
+        x9 = self.decoder_block3(x8, x4)  # (bs, 256, T / 8, F / 8)
+        x10 = self.decoder_block4(x9, x3)  # (bs, 128, T / 4, F / 4)
+        x11 = self.decoder_block5(x10, x2)  # (bs, 64, T / 2, F / 2)
+        x12 = self.decoder_block6(x11, x1)  # (bs, 32, T, F)
+        (x, _) = self.after_conv_block1(x12)  # (bs, 32, T, F)
+
+        x = self.after_conv2(x)  # (bs, channels * 3, T, F)
+        # (batch_size, input_channles * subbands_num * targets_num * k, T, F')
+
+        x = self.subband.synthesis(x)
+        # (batch_size, input_channles * targets_num * K, T, F)
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 1024 -> 1025.
+        x = x[:, :, 0:origin_len, :]  # (bs, feature_maps, time_steps, freq_bins)
+
+        audio_length = mixtures.shape[2]
+
+        separated_audio = self.feature_maps_to_wav(x, mag, sin_in, cos_in, audio_length)
+        # separated_audio: (batch_size, target_sources_num * input_channels, segments_num)
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/models/resunet_ismir2021.py

@@ -0,0 +1,534 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from inplace_abn.abn import InPlaceABNSync
+from torchlibrosa.stft import ISTFT, STFT, magphase
+
+from bytesep.models.pytorch_modules import Base, init_bn, init_layer
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, activation, momentum):
+        r"""Residual block."""
+        super(ConvBlockRes, self).__init__()
+
+        self.activation = activation
+        padding = [kernel_size[0] // 2, kernel_size[1] // 2]
+
+        # ABN is not used for bn1 because we found using abn1 will degrade performance.
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+
+        self.abn2 = InPlaceABNSync(
+            num_features=out_channels, momentum=momentum, activation='leaky_relu'
+        )
+
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(1, 1),
+                stride=(1, 1),
+                padding=(0, 0),
+            )
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+
+        if self.is_shortcut:
+            init_layer(self.shortcut)
+
+    def forward(self, x):
+        origin = x
+        x = self.conv1(F.leaky_relu_(self.bn1(x), negative_slope=0.01))
+        x = self.conv2(self.abn2(x))
+
+        if self.is_shortcut:
+            return self.shortcut(origin) + x
+        else:
+            return origin + x
+
+
+class EncoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, downsample, activation, momentum
+    ):
+        r"""Encoder block, contains 8 convolutional layers."""
+        super(EncoderBlockRes4B, self).__init__()
+
+        self.conv_block1 = ConvBlockRes(
+            in_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block2 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.downsample = downsample
+
+    def forward(self, x):
+        encoder = self.conv_block1(x)
+        encoder = self.conv_block2(encoder)
+        encoder = self.conv_block3(encoder)
+        encoder = self.conv_block4(encoder)
+        encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample)
+        return encoder_pool, encoder
+
+
+class DecoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, upsample, activation, momentum
+    ):
+        r"""Decoder block, contains 1 transpose convolutional and 8 convolutional layers."""
+        super(DecoderBlockRes4B, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = upsample
+        self.activation = activation
+
+        self.conv1 = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=self.stride,
+            stride=self.stride,
+            padding=(0, 0),
+            bias=False,
+            dilation=(1, 1),
+        )
+
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.conv_block2 = ConvBlockRes(
+            out_channels * 2, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block5 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_layer(self.conv1)
+
+    def forward(self, input_tensor, concat_tensor):
+        x = self.conv1(F.relu_(self.bn1(input_tensor)))
+        x = torch.cat((x, concat_tensor), dim=1)
+        x = self.conv_block2(x)
+        x = self.conv_block3(x)
+        x = self.conv_block4(x)
+        x = self.conv_block5(x)
+        return x
+
+
+class ResUNet143_DecouplePlusInplaceABN_ISMIR2021(nn.Module, Base):
+    def __init__(self, input_channels, target_sources_num):
+        super(ResUNet143_DecouplePlusInplaceABN_ISMIR2021, self).__init__()
+
+        self.input_channels = input_channels
+        self.target_sources_num = target_sources_num
+
+        window_size = 2048
+        hop_size = 441
+        center = True
+        pad_mode = 'reflect'
+        window = 'hann'
+        activation = 'leaky_relu'
+        momentum = 0.01
+
+        self.subbands_num = 1
+
+        assert (
+            self.subbands_num == 1
+        ), "Using subbands_num > 1 on spectrogram \
+            will lead to unexpected performance sometimes. Suggest to use \
+            subband method on waveform."
+
+        # Downsample rate along the time axis.
+        self.K = 4  # outputs: |M|, cos∠M, sin∠M, Q
+        self.time_downsample_ratio = 2 ** 5  # This number equals 2^{#encoder_blcoks}
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.encoder_block1 = EncoderBlockRes4B(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlockRes4B(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlockRes4B(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlockRes4B(
+            in_channels=256,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7a = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7b = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7c = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7d = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=256,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlockRes4B(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlockRes4B(
+            in_channels=128,
+            out_channels=64,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block6 = DecoderBlockRes4B(
+            in_channels=64,
+            out_channels=32,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=target_sources_num
+            * input_channels
+            * self.K
+            * self.subbands_num,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(
+        self,
+        input_tensor: torch.Tensor,
+        sp: torch.Tensor,
+        sin_in: torch.Tensor,
+        cos_in: torch.Tensor,
+        audio_length: int,
+    ) -> torch.Tensor:
+        r"""Convert feature maps to waveform.
+
+        Args:
+            input_tensor: (batch_size, target_sources_num * input_channels * self.K, time_steps, freq_bins)
+            sp: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            sin_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            cos_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+
+        Outputs:
+            waveform: (batch_size, target_sources_num * input_channels, segment_samples)
+        """
+        batch_size, _, time_steps, freq_bins = input_tensor.shape
+
+        x = input_tensor.reshape(
+            batch_size,
+            self.target_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, target_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        linear_mag = x[:, :, :, 3, :, :]
+        # mask_cos, mask_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+        # out_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag + linear_mag)
+        # out_mag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.target_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        x = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * target_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        waveform = x.reshape(
+            batch_size, self.target_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, target_sources_num * input_channels, segments_num)
+
+        return waveform
+
+    def forward(self, input_dict):
+        r"""Forward data into the module.
+
+        Args:
+            input_dict: dict, e.g., {
+                waveform: (batch_size, input_channels, segment_samples),
+                ...,
+            }
+
+        Outputs:
+            output_dict: dict, e.g., {
+                'waveform': (batch_size, input_channels, segment_samples),
+                ...,
+            }
+        """
+        mixtures = input_dict['waveform']
+        # (batch_size, input_channels, segment_samples)
+
+        mag, cos_in, sin_in = self.wav_to_spectrogram_phase(mixtures)
+        # mag, cos_in, sin_in: (batch_size, input_channels, time_steps, freq_bins)
+
+        # Batch normalize on individual frequency bins.
+        x = mag.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        # x: (batch_size, input_channels, time_steps, freq_bins)
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.time_downsample_ratio))
+            * self.time_downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        # (batch_size, channels, padded_time_steps, freq_bins)
+
+        # Let frequency bins be evenly divided by 2, e.g., 1025 -> 1024.
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, channels, T, F)
+
+        if self.subbands_num > 1:
+            x = self.subband.analysis(x)
+            # (bs, input_channels, T, F'), where F' = F // subbands_num
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F / 2)
+        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F / 4)
+        (x3_pool, x3) = self.encoder_block3(x2_pool)  # x3_pool: (bs, 128, T / 8, F / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool
+        )  # x4_pool: (bs, 256, T / 16, F / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool
+        )  # x5_pool: (bs, 384, T / 32, F / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool
+        )  # x6_pool: (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7a(x6_pool)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7b(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7c(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7d(x_center)  # (bs, 384, T / 32, F / 64)
+        x7 = self.decoder_block1(x_center, x6)  # (bs, 384, T / 32, F / 32)
+        x8 = self.decoder_block2(x7, x5)  # (bs, 384, T / 16, F / 16)
+        x9 = self.decoder_block3(x8, x4)  # (bs, 256, T / 8, F / 8)
+        x10 = self.decoder_block4(x9, x3)  # (bs, 128, T / 4, F / 4)
+        x11 = self.decoder_block5(x10, x2)  # (bs, 64, T / 2, F / 2)
+        x12 = self.decoder_block6(x11, x1)  # (bs, 32, T, F)
+        (x, _) = self.after_conv_block1(x12)  # (bs, 32, T, F)
+
+        x = self.after_conv2(x)  # (bs, channels * 3, T, F)
+        # (batch_size, target_sources_num * input_channles * self.K * subbands_num, T, F')
+
+        if self.subbands_num > 1:
+            x = self.subband.synthesis(x)
+            # (batch_size, target_sources_num * input_channles * self.K, T, F)
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 1024 -> 1025.
+
+        x = x[:, :, 0:origin_len, :]
+        # (batch_size, target_sources_num * input_channles * self.K, T, F)
+
+        audio_length = mixtures.shape[2]
+
+        separated_audio = self.feature_maps_to_wav(x, mag, sin_in, cos_in, audio_length)
+        # separated_audio: (batch_size, target_sources_num * input_channels, segments_num)
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/models/resunet_subbandtime.py

@@ -0,0 +1,545 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchlibrosa.stft import ISTFT, STFT, magphase
+
+from bytesep.models.pytorch_modules import Base, init_bn, init_layer
+from bytesep.models.subband_tools.pqmf import PQMF
+
+
+class ConvBlockRes(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, activation, momentum):
+        r"""Residual block."""
+        super(ConvBlockRes, self).__init__()
+
+        self.activation = activation
+        padding = [kernel_size[0] // 2, kernel_size[1] // 2]
+
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(1, 1),
+                stride=(1, 1),
+                padding=(0, 0),
+            )
+            self.is_shortcut = True
+        else:
+            self.is_shortcut = False
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+
+        if self.is_shortcut:
+            init_layer(self.shortcut)
+
+    def forward(self, x):
+        origin = x
+        x = self.conv1(F.leaky_relu_(self.bn1(x), negative_slope=0.01))
+        x = self.conv2(F.leaky_relu_(self.bn2(x), negative_slope=0.01))
+
+        if self.is_shortcut:
+            return self.shortcut(origin) + x
+        else:
+            return origin + x
+
+
+class EncoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, downsample, activation, momentum
+    ):
+        r"""Encoder block, contains 8 convolutional layers."""
+        super(EncoderBlockRes4B, self).__init__()
+
+        self.conv_block1 = ConvBlockRes(
+            in_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block2 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.downsample = downsample
+
+    def forward(self, x):
+        encoder = self.conv_block1(x)
+        encoder = self.conv_block2(encoder)
+        encoder = self.conv_block3(encoder)
+        encoder = self.conv_block4(encoder)
+        encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample)
+        return encoder_pool, encoder
+
+
+class DecoderBlockRes4B(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, upsample, activation, momentum
+    ):
+        r"""Decoder block, contains 1 transpose convolutional and 8 convolutional layers."""
+        super(DecoderBlockRes4B, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = upsample
+        self.activation = activation
+
+        self.conv1 = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=self.stride,
+            stride=self.stride,
+            padding=(0, 0),
+            bias=False,
+            dilation=(1, 1),
+        )
+
+        self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum)
+        self.conv_block2 = ConvBlockRes(
+            out_channels * 2, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block3 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block4 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.conv_block5 = ConvBlockRes(
+            out_channels, out_channels, kernel_size, activation, momentum
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn1)
+        init_layer(self.conv1)
+
+    def forward(self, input_tensor, concat_tensor):
+        x = self.conv1(F.relu_(self.bn1(input_tensor)))
+        x = torch.cat((x, concat_tensor), dim=1)
+        x = self.conv_block2(x)
+        x = self.conv_block3(x)
+        x = self.conv_block4(x)
+        x = self.conv_block5(x)
+        return x
+
+
+class ResUNet143_Subbandtime(nn.Module, Base):
+    def __init__(self, input_channels, target_sources_num):
+        super(ResUNet143_Subbandtime, self).__init__()
+
+        self.input_channels = input_channels
+        self.target_sources_num = target_sources_num
+
+        window_size = 512
+        hop_size = 110
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+        activation = "leaky_relu"
+        momentum = 0.01
+
+        self.subbands_num = 4
+        self.K = 4  # outputs: |M|, cos∠M, sin∠M, Q
+
+        self.downsample_ratio = 2 ** 5  # This number equals 2^{#encoder_blcoks}
+
+        self.pqmf = PQMF(
+            N=self.subbands_num,
+            M=64,
+            project_root='bytesep/models/subband_tools/filters',
+        )
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.encoder_block1 = EncoderBlockRes4B(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlockRes4B(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlockRes4B(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlockRes4B(
+            in_channels=256,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7a = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7b = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7c = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7d = EncoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(1, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlockRes4B(
+            in_channels=384,
+            out_channels=256,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlockRes4B(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlockRes4B(
+            in_channels=128,
+            out_channels=64,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block6 = DecoderBlockRes4B(
+            in_channels=64,
+            out_channels=32,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = EncoderBlockRes4B(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(1, 1),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=target_sources_num
+            * input_channels
+            * self.K
+            * self.subbands_num,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(
+        self,
+        input_tensor: torch.Tensor,
+        sp: torch.Tensor,
+        sin_in: torch.Tensor,
+        cos_in: torch.Tensor,
+        audio_length: int,
+    ) -> torch.Tensor:
+        r"""Convert feature maps to waveform.
+
+        Args:
+            input_tensor: (batch_size, target_sources_num * input_channels * self.K, time_steps, freq_bins)
+            sp: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            sin_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            cos_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+
+        Outputs:
+            waveform: (batch_size, target_sources_num * input_channels, segment_samples)
+        """
+        batch_size, _, time_steps, freq_bins = input_tensor.shape
+
+        x = input_tensor.reshape(
+            batch_size,
+            self.target_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, target_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        linear_mag = torch.tanh(x[:, :, :, 3, :, :])
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        # mask_cos, mask_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+        # out_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag + linear_mag)
+        # out_mag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.target_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        x = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * target_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        waveform = x.reshape(
+            batch_size, self.target_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, target_sources_num * input_channels, segments_num)
+
+        return waveform
+
+    def forward(self, input_dict):
+        r"""Forward data into the module.
+
+        Args:
+            input_dict: dict, e.g., {
+                waveform: (batch_size, input_channels, segment_samples),
+                ...,
+            }
+
+        Outputs:
+            output_dict: dict, e.g., {
+                'waveform': (batch_size, input_channels, segment_samples),
+                ...,
+            }
+        """
+        mixtures = input_dict['waveform']
+        # (batch_size, input_channels, segment_samples)
+
+        subband_x = self.pqmf.analysis(mixtures)
+        # subband_x: (batch_size, input_channels * subbands_num, segment_samples)
+
+        mag, cos_in, sin_in = self.wav_to_spectrogram_phase(subband_x)
+        # mag, cos_in, sin_in: (batch_size, input_channels * subbands_num, time_steps, freq_bins)
+
+        # Batch normalize on individual frequency bins.
+        x = mag.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        # (batch_size, input_channels * subbands_num, time_steps, freq_bins)
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        # x: (batch_size, input_channels * subbands_num, padded_time_steps, freq_bins)
+
+        # Let frequency bins be evenly divided by 2, e.g., 257 -> 256
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, input_channels, T, F)
+        # x: (batch_size, input_channels * subbands_num, padded_time_steps, freq_bins)
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F / 2)
+        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F / 4)
+        (x3_pool, x3) = self.encoder_block3(x2_pool)  # x3_pool: (bs, 128, T / 8, F / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool
+        )  # x4_pool: (bs, 256, T / 16, F / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool
+        )  # x5_pool: (bs, 384, T / 32, F / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool
+        )  # x6_pool: (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7a(x6_pool)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7b(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7c(x_center)  # (bs, 384, T / 32, F / 64)
+        (x_center, _) = self.conv_block7d(x_center)  # (bs, 384, T / 32, F / 64)
+        x7 = self.decoder_block1(x_center, x6)  # (bs, 384, T / 32, F / 32)
+        x8 = self.decoder_block2(x7, x5)  # (bs, 384, T / 16, F / 16)
+        x9 = self.decoder_block3(x8, x4)  # (bs, 256, T / 8, F / 8)
+        x10 = self.decoder_block4(x9, x3)  # (bs, 128, T / 4, F / 4)
+        x11 = self.decoder_block5(x10, x2)  # (bs, 64, T / 2, F / 2)
+        x12 = self.decoder_block6(x11, x1)  # (bs, 32, T, F)
+        (x, _) = self.after_conv_block1(x12)  # (bs, 32, T, F)
+
+        x = self.after_conv2(x)
+        # (batch_size, subbands_num * target_sources_num * input_channles * self.K, T, F')
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 256 -> 257.
+
+        x = x[:, :, 0:origin_len, :]
+        # (batch_size, subbands_num * target_sources_num * input_channles * self.K, T, F')
+
+        audio_length = subband_x.shape[2]
+
+        # Recover each subband spectrograms to subband waveforms. Then synthesis
+        # the subband waveforms to a waveform.
+        C1 = x.shape[1] // self.subbands_num
+        C2 = mag.shape[1] // self.subbands_num
+
+        separated_subband_audio = torch.cat(
+            [
+                self.feature_maps_to_wav(
+                    input_tensor=x[:, j * C1 : (j + 1) * C1, :, :],
+                    sp=mag[:, j * C2 : (j + 1) * C2, :, :],
+                    sin_in=sin_in[:, j * C2 : (j + 1) * C2, :, :],
+                    cos_in=cos_in[:, j * C2 : (j + 1) * C2, :, :],
+                    audio_length=audio_length,
+                )
+                for j in range(self.subbands_num)
+            ],
+            dim=1,
+        )
+        # （batch_size, subbands_num * target_sources_num * input_channles, segment_samples)
+
+        separated_audio = self.pqmf.synthesis(separated_subband_audio)
+        # (batch_size, input_channles, segment_samples)
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/models/subband_tools/fDomainHelper.py

@@ -0,0 +1,255 @@
+from torchlibrosa.stft import STFT, ISTFT, magphase
+import torch
+import torch.nn as nn
+import numpy as np
+from tools.pytorch.modules.pqmf import PQMF
+
+
+class FDomainHelper(nn.Module):
+    def __init__(
+        self,
+        window_size=2048,
+        hop_size=441,
+        center=True,
+        pad_mode='reflect',
+        window='hann',
+        freeze_parameters=True,
+        subband=None,
+        root="/Users/admin/Documents/projects/",
+    ):
+        super(FDomainHelper, self).__init__()
+        self.subband = subband
+        if self.subband is None:
+            self.stft = STFT(
+                n_fft=window_size,
+                hop_length=hop_size,
+                win_length=window_size,
+                window=window,
+                center=center,
+                pad_mode=pad_mode,
+                freeze_parameters=freeze_parameters,
+            )
+
+            self.istft = ISTFT(
+                n_fft=window_size,
+                hop_length=hop_size,
+                win_length=window_size,
+                window=window,
+                center=center,
+                pad_mode=pad_mode,
+                freeze_parameters=freeze_parameters,
+            )
+        else:
+            self.stft = STFT(
+                n_fft=window_size // self.subband,
+                hop_length=hop_size // self.subband,
+                win_length=window_size // self.subband,
+                window=window,
+                center=center,
+                pad_mode=pad_mode,
+                freeze_parameters=freeze_parameters,
+            )
+
+            self.istft = ISTFT(
+                n_fft=window_size // self.subband,
+                hop_length=hop_size // self.subband,
+                win_length=window_size // self.subband,
+                window=window,
+                center=center,
+                pad_mode=pad_mode,
+                freeze_parameters=freeze_parameters,
+            )
+
+        if subband is not None and root is not None:
+            self.qmf = PQMF(subband, 64, root)
+
+    def complex_spectrogram(self, input, eps=0.0):
+        # [batchsize, samples]
+        # return [batchsize, 2, t-steps, f-bins]
+        real, imag = self.stft(input)
+        return torch.cat([real, imag], dim=1)
+
+    def reverse_complex_spectrogram(self, input, eps=0.0, length=None):
+        # [batchsize, 2[real,imag], t-steps, f-bins]
+        wav = self.istft(input[:, 0:1, ...], input[:, 1:2, ...], length=length)
+        return wav
+
+    def spectrogram(self, input, eps=0.0):
+        (real, imag) = self.stft(input.float())
+        return torch.clamp(real ** 2 + imag ** 2, eps, np.inf) ** 0.5
+
+    def spectrogram_phase(self, input, eps=0.0):
+        (real, imag) = self.stft(input.float())
+        mag = torch.clamp(real ** 2 + imag ** 2, eps, np.inf) ** 0.5
+        cos = real / mag
+        sin = imag / mag
+        return mag, cos, sin
+
+    def wav_to_spectrogram_phase(self, input, eps=1e-8):
+        """Waveform to spectrogram.
+
+        Args:
+          input: (batch_size, channels_num, segment_samples)
+
+        Outputs:
+          output: (batch_size, channels_num, time_steps, freq_bins)
+        """
+        sp_list = []
+        cos_list = []
+        sin_list = []
+        channels_num = input.shape[1]
+        for channel in range(channels_num):
+            mag, cos, sin = self.spectrogram_phase(input[:, channel, :], eps=eps)
+            sp_list.append(mag)
+            cos_list.append(cos)
+            sin_list.append(sin)
+
+        sps = torch.cat(sp_list, dim=1)
+        coss = torch.cat(cos_list, dim=1)
+        sins = torch.cat(sin_list, dim=1)
+        return sps, coss, sins
+
+    def spectrogram_phase_to_wav(self, sps, coss, sins, length):
+        channels_num = sps.size()[1]
+        res = []
+        for i in range(channels_num):
+            res.append(
+                self.istft(
+                    sps[:, i : i + 1, ...] * coss[:, i : i + 1, ...],
+                    sps[:, i : i + 1, ...] * sins[:, i : i + 1, ...],
+                    length,
+                )
+            )
+            res[-1] = res[-1].unsqueeze(1)
+        return torch.cat(res, dim=1)
+
+    def wav_to_spectrogram(self, input, eps=1e-8):
+        """Waveform to spectrogram.
+
+        Args:
+          input: (batch_size,channels_num, segment_samples)
+
+        Outputs:
+          output: (batch_size, channels_num, time_steps, freq_bins)
+        """
+        sp_list = []
+        channels_num = input.shape[1]
+        for channel in range(channels_num):
+            sp_list.append(self.spectrogram(input[:, channel, :], eps=eps))
+        output = torch.cat(sp_list, dim=1)
+        return output
+
+    def spectrogram_to_wav(self, input, spectrogram, length=None):
+        """Spectrogram to waveform.
+        Args:
+          input: (batch_size, segment_samples, channels_num)
+          spectrogram: (batch_size, channels_num, time_steps, freq_bins)
+
+        Outputs:
+          output: (batch_size, segment_samples, channels_num)
+        """
+        channels_num = input.shape[1]
+        wav_list = []
+        for channel in range(channels_num):
+            (real, imag) = self.stft(input[:, channel, :])
+            (_, cos, sin) = magphase(real, imag)
+            wav_list.append(
+                self.istft(
+                    spectrogram[:, channel : channel + 1, :, :] * cos,
+                    spectrogram[:, channel : channel + 1, :, :] * sin,
+                    length,
+                )
+            )
+
+        output = torch.stack(wav_list, dim=1)
+        return output
+
+    # todo the following code is not bug free!
+    def wav_to_complex_spectrogram(self, input, eps=0.0):
+        # [batchsize , channels, samples]
+        # [batchsize, 2[real,imag]*channels, t-steps, f-bins]
+        res = []
+        channels_num = input.shape[1]
+        for channel in range(channels_num):
+            res.append(self.complex_spectrogram(input[:, channel, :], eps=eps))
+        return torch.cat(res, dim=1)
+
+    def complex_spectrogram_to_wav(self, input, eps=0.0, length=None):
+        # [batchsize, 2[real,imag]*channels, t-steps, f-bins]
+        # return  [batchsize, channels, samples]
+        channels = input.size()[1] // 2
+        wavs = []
+        for i in range(channels):
+            wavs.append(
+                self.reverse_complex_spectrogram(
+                    input[:, 2 * i : 2 * i + 2, ...], eps=eps, length=length
+                )
+            )
+            wavs[-1] = wavs[-1].unsqueeze(1)
+        return torch.cat(wavs, dim=1)
+
+    def wav_to_complex_subband_spectrogram(self, input, eps=0.0):
+        # [batchsize, channels, samples]
+        # [batchsize, 2[real,imag]*subband*channels, t-steps, f-bins]
+        subwav = self.qmf.analysis(input)  # [batchsize, subband*channels, samples]
+        subspec = self.wav_to_complex_spectrogram(subwav)
+        return subspec
+
+    def complex_subband_spectrogram_to_wav(self, input, eps=0.0):
+        # [batchsize, 2[real,imag]*subband*channels, t-steps, f-bins]
+        # [batchsize, channels, samples]
+        subwav = self.complex_spectrogram_to_wav(input)
+        data = self.qmf.synthesis(subwav)
+        return data
+
+    def wav_to_mag_phase_subband_spectrogram(self, input, eps=1e-8):
+        """
+        :param input:
+        :param eps:
+        :return:
+            loss = torch.nn.L1Loss()
+            model = FDomainHelper(subband=4)
+            data = torch.randn((3,1, 44100*3))
+
+            sps, coss, sins = model.wav_to_mag_phase_subband_spectrogram(data)
+            wav = model.mag_phase_subband_spectrogram_to_wav(sps,coss,sins,44100*3//4)
+
+            print(loss(data,wav))
+            print(torch.max(torch.abs(data-wav)))
+
+        """
+        # [batchsize, channels, samples]
+        # [batchsize, 2[real,imag]*subband*channels, t-steps, f-bins]
+        subwav = self.qmf.analysis(input)  # [batchsize, subband*channels, samples]
+        sps, coss, sins = self.wav_to_spectrogram_phase(subwav, eps=eps)
+        return sps, coss, sins
+
+    def mag_phase_subband_spectrogram_to_wav(self, sps, coss, sins, length, eps=0.0):
+        # [batchsize, 2[real,imag]*subband*channels, t-steps, f-bins]
+        # [batchsize, channels, samples]
+        subwav = self.spectrogram_phase_to_wav(sps, coss, sins, length)
+        data = self.qmf.synthesis(subwav)
+        return data
+
+
+if __name__ == "__main__":
+    # from thop import profile
+    # from thop import clever_format
+    # from tools.file.wav import *
+    # import time
+    #
+    # wav = torch.randn((1,2,44100))
+    # model = FDomainHelper()
+
+    from tools.file.wav import *
+
+    loss = torch.nn.L1Loss()
+    model = FDomainHelper()
+    data = torch.randn((3, 1, 44100 * 5))
+
+    sps = model.wav_to_complex_spectrogram(data)
+    print(sps.size())
+    wav = model.complex_spectrogram_to_wav(sps, 44100 * 5)
+
+    print(loss(data, wav))
+    print(torch.max(torch.abs(data - wav)))

bytesep/models/subband_tools/pqmf.py

@@ -0,0 +1,136 @@
+'''
+@File    :   subband_util.py
+@Contact :   [email protected]
+@License :   (C)Copyright 2020-2021
+@Modify Time      @Author    @Version    @Desciption
+------------      -------    --------    -----------
+2020/4/3 4:54 PM   Haohe Liu      1.0         None
+'''
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import numpy as np
+import os.path as op
+from scipy.io import loadmat
+
+
+def load_mat2numpy(fname=""):
+    '''
+    Args:
+        fname: pth to mat
+        type:
+    Returns: dic object
+    '''
+    if len(fname) == 0:
+        return None
+    else:
+        return loadmat(fname)
+
+
+class PQMF(nn.Module):
+    def __init__(self, N, M, project_root):
+        super().__init__()
+        self.N = N  # nsubband
+        self.M = M  # nfilter
+        try:
+            assert (N, M) in [(8, 64), (4, 64), (2, 64)]
+        except:
+            print("Warning:", N, "subbandand ", M, " filter is not supported")
+        self.pad_samples = 64
+        self.name = str(N) + "_" + str(M) + ".mat"
+        self.ana_conv_filter = nn.Conv1d(
+            1, out_channels=N, kernel_size=M, stride=N, bias=False
+        )
+        data = load_mat2numpy(op.join(project_root, "f_" + self.name))
+        data = data['f'].astype(np.float32) / N
+        data = np.flipud(data.T).T
+        data = np.reshape(data, (N, 1, M)).copy()
+        dict_new = self.ana_conv_filter.state_dict().copy()
+        dict_new['weight'] = torch.from_numpy(data)
+        self.ana_pad = nn.ConstantPad1d((M - N, 0), 0)
+        self.ana_conv_filter.load_state_dict(dict_new)
+
+        self.syn_pad = nn.ConstantPad1d((0, M // N - 1), 0)
+        self.syn_conv_filter = nn.Conv1d(
+            N, out_channels=N, kernel_size=M // N, stride=1, bias=False
+        )
+        gk = load_mat2numpy(op.join(project_root, "h_" + self.name))
+        gk = gk['h'].astype(np.float32)
+        gk = np.transpose(np.reshape(gk, (N, M // N, N)), (1, 0, 2)) * N
+        gk = np.transpose(gk[::-1, :, :], (2, 1, 0)).copy()
+        dict_new = self.syn_conv_filter.state_dict().copy()
+        dict_new['weight'] = torch.from_numpy(gk)
+        self.syn_conv_filter.load_state_dict(dict_new)
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def __analysis_channel(self, inputs):
+        return self.ana_conv_filter(self.ana_pad(inputs))
+
+    def __systhesis_channel(self, inputs):
+        ret = self.syn_conv_filter(self.syn_pad(inputs)).permute(0, 2, 1)
+        return torch.reshape(ret, (ret.shape[0], 1, -1))
+
+    def analysis(self, inputs):
+        '''
+        :param inputs: [batchsize,channel,raw_wav],value:[0,1]
+        :return:
+        '''
+        inputs = F.pad(inputs, ((0, self.pad_samples)))
+        ret = None
+        for i in range(inputs.size()[1]):  # channels
+            if ret is None:
+                ret = self.__analysis_channel(inputs[:, i : i + 1, :])
+            else:
+                ret = torch.cat(
+                    (ret, self.__analysis_channel(inputs[:, i : i + 1, :])), dim=1
+                )
+        return ret
+
+    def synthesis(self, data):
+        '''
+        :param data: [batchsize,self.N*K,raw_wav_sub],value:[0,1]
+        :return:
+        '''
+        ret = None
+        # data = F.pad(data,((0,self.pad_samples//self.N)))
+        for i in range(data.size()[1]):  # channels
+            if i % self.N == 0:
+                if ret is None:
+                    ret = self.__systhesis_channel(data[:, i : i + self.N, :])
+                else:
+                    new = self.__systhesis_channel(data[:, i : i + self.N, :])
+                    ret = torch.cat((ret, new), dim=1)
+        ret = ret[..., : -self.pad_samples]
+        return ret
+
+    def forward(self, inputs):
+        return self.ana_conv_filter(self.ana_pad(inputs))
+
+
+if __name__ == "__main__":
+    import torch
+    import numpy as np
+    import matplotlib.pyplot as plt
+    from tools.file.wav import *
+
+    pqmf = PQMF(N=4, M=64, project_root="/Users/admin/Documents/projects")
+
+    rs = np.random.RandomState(0)
+    x = torch.tensor(rs.rand(4, 2, 32000), dtype=torch.float32)
+
+    a1 = pqmf.analysis(x)
+    a2 = pqmf.synthesis(a1)
+
+    print(a2.size(), x.size())
+
+    plt.subplot(211)
+    plt.plot(x[0, 0, -500:])
+    plt.subplot(212)
+    plt.plot(a2[0, 0, -500:])
+    plt.plot(x[0, 0, -500:] - a2[0, 0, -500:])
+    plt.show()
+
+    print(torch.sum(torch.abs(x[...] - a2[...])))

bytesep/models/unet.py

@@ -0,0 +1,532 @@
+import math
+from typing import Dict, List, NoReturn, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.optim.lr_scheduler import LambdaLR
+from torchlibrosa.stft import ISTFT, STFT, magphase
+
+from bytesep.models.pytorch_modules import Base, Subband, act, init_bn, init_layer
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Tuple,
+        activation: str,
+        momentum: float,
+    ):
+        r"""Convolutional block."""
+        super(ConvBlock, self).__init__()
+
+        self.activation = activation
+        padding = (kernel_size[0] // 2, kernel_size[1] // 2)
+
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv2 = nn.Conv2d(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=(1, 1),
+            dilation=(1, 1),
+            padding=padding,
+            bias=False,
+        )
+
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.init_weights()
+
+    def init_weights(self) -> NoReturn:
+        r"""Initialize weights."""
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        r"""Forward data into the module.
+
+        Args:
+            input_tensor: (batch_size, in_feature_maps, time_steps, freq_bins)
+
+        Returns:
+            output_tensor: (batch_size, out_feature_maps, time_steps, freq_bins)
+        """
+        x = act(self.bn1(self.conv1(input_tensor)), self.activation)
+        x = act(self.bn2(self.conv2(x)), self.activation)
+        output_tensor = x
+
+        return output_tensor
+
+
+class EncoderBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Tuple,
+        downsample: Tuple,
+        activation: str,
+        momentum: float,
+    ):
+        r"""Encoder block."""
+        super(EncoderBlock, self).__init__()
+
+        self.conv_block = ConvBlock(
+            in_channels, out_channels, kernel_size, activation, momentum
+        )
+        self.downsample = downsample
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        r"""Forward data into the module.
+
+        Args:
+            input_tensor: (batch_size, in_feature_maps, time_steps, freq_bins)
+
+        Returns:
+            encoder_pool: (batch_size, out_feature_maps, downsampled_time_steps, downsampled_freq_bins)
+            encoder: (batch_size, out_feature_maps, time_steps, freq_bins)
+        """
+        encoder_tensor = self.conv_block(input_tensor)
+        # encoder: (batch_size, out_feature_maps, time_steps, freq_bins)
+
+        encoder_pool = F.avg_pool2d(encoder_tensor, kernel_size=self.downsample)
+        # encoder_pool: (batch_size, out_feature_maps, downsampled_time_steps, downsampled_freq_bins)
+
+        return encoder_pool, encoder_tensor
+
+
+class DecoderBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Tuple,
+        upsample: Tuple,
+        activation: str,
+        momentum: float,
+    ):
+        r"""Decoder block."""
+        super(DecoderBlock, self).__init__()
+
+        self.kernel_size = kernel_size
+        self.stride = upsample
+        self.activation = activation
+
+        self.conv1 = torch.nn.ConvTranspose2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=self.stride,
+            stride=self.stride,
+            padding=(0, 0),
+            bias=False,
+            dilation=(1, 1),
+        )
+
+        self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)
+
+        self.conv_block2 = ConvBlock(
+            out_channels * 2, out_channels, kernel_size, activation, momentum
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        r"""Initialize weights."""
+        init_layer(self.conv1)
+        init_bn(self.bn1)
+
+    def forward(
+        self, input_tensor: torch.Tensor, concat_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        r"""Forward data into the module.
+
+        Args:
+            torch_tensor: (batch_size, in_feature_maps, downsampled_time_steps, downsampled_freq_bins)
+            concat_tensor: (batch_size, in_feature_maps, time_steps, freq_bins)
+
+        Returns:
+            output_tensor: (batch_size, out_feature_maps, time_steps, freq_bins)
+        """
+        x = act(self.bn1(self.conv1(input_tensor)), self.activation)
+        # (batch_size, in_feature_maps, time_steps, freq_bins)
+
+        x = torch.cat((x, concat_tensor), dim=1)
+        # (batch_size, in_feature_maps * 2, time_steps, freq_bins)
+
+        output_tensor = self.conv_block2(x)
+        # output_tensor: (batch_size, out_feature_maps, time_steps, freq_bins)
+
+        return output_tensor
+
+
+class UNet(nn.Module, Base):
+    def __init__(self, input_channels: int, target_sources_num: int):
+        r"""UNet."""
+        super(UNet, self).__init__()
+
+        self.input_channels = input_channels
+        self.target_sources_num = target_sources_num
+
+        window_size = 2048
+        hop_size = 441
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+        activation = "leaky_relu"
+        momentum = 0.01
+
+        self.subbands_num = 1
+
+        assert (
+            self.subbands_num == 1
+        ), "Using subbands_num > 1 on spectrogram \
+            will lead to unexpected performance sometimes. Suggest to use \
+            subband method on waveform."
+
+        self.K = 3  # outputs: |M|, cos∠M, sin∠M
+        self.downsample_ratio = 2 ** 6  # This number equals 2^{#encoder_blcoks}
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.subband = Subband(subbands_num=self.subbands_num)
+
+        self.encoder_block1 = EncoderBlock(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlock(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlock(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlock(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlock(
+            in_channels=256,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7 = ConvBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlock(
+            in_channels=384,
+            out_channels=256,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlock(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlock(
+            in_channels=128,
+            out_channels=64,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.decoder_block6 = DecoderBlock(
+            in_channels=64,
+            out_channels=32,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = ConvBlock(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=target_sources_num
+            * input_channels
+            * self.K
+            * self.subbands_num,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        r"""Initialize weights."""
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(
+        self,
+        input_tensor: torch.Tensor,
+        sp: torch.Tensor,
+        sin_in: torch.Tensor,
+        cos_in: torch.Tensor,
+        audio_length: int,
+    ) -> torch.Tensor:
+        r"""Convert feature maps to waveform.
+
+        Args:
+            input_tensor: (batch_size, target_sources_num * input_channels * self.K, time_steps, freq_bins)
+            sp: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            sin_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            cos_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+
+        Outputs:
+            waveform: (batch_size, target_sources_num * input_channels, segment_samples)
+        """
+        batch_size, _, time_steps, freq_bins = input_tensor.shape
+
+        x = input_tensor.reshape(
+            batch_size,
+            self.target_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, target_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        # mask_cos, mask_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+        # out_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag)
+        # out_mag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.target_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        x = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * target_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        waveform = x.reshape(
+            batch_size, self.target_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, target_sources_num * input_channels, segments_num)
+
+        return waveform
+
+    def forward(self, input_dict: Dict) -> Dict:
+        r"""Forward data into the module.
+
+        Args:
+            input_dict: dict, e.g., {
+                waveform: (batch_size, input_channels, segment_samples),
+                ...,
+            }
+
+        Outputs:
+            output_dict: dict, e.g., {
+                'waveform': (batch_size, input_channels, segment_samples),
+                ...,
+            }
+        """
+        mixtures = input_dict['waveform']
+        # (batch_size, input_channels, segment_samples)
+
+        mag, cos_in, sin_in = self.wav_to_spectrogram_phase(mixtures)
+        # mag, cos_in, sin_in: (batch_size, input_channels, time_steps, freq_bins)
+
+        # Batch normalize on individual frequency bins.
+        x = mag.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        # x: (batch_size, input_channels, time_steps, freq_bins)
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        # x: (batch_size, input_channels, padded_time_steps, freq_bins)
+
+        # Let frequency bins be evenly divided by 2, e.g., 1025 -> 1024
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, input_channels, T, F)
+
+        if self.subbands_num > 1:
+            x = self.subband.analysis(x)
+            # (bs, input_channels, T, F'), where F' = F // subbands_num
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F' / 2)
+        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F' / 4)
+        (x3_pool, x3) = self.encoder_block3(
+            x2_pool
+        )  # x3_pool: (bs, 128, T / 8, F' / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool
+        )  # x4_pool: (bs, 256, T / 16, F' / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool
+        )  # x5_pool: (bs, 384, T / 32, F' / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool
+        )  # x6_pool: (bs, 384, T / 64, F' / 64)
+        x_center = self.conv_block7(x6_pool)  # (bs, 384, T / 64, F' / 64)
+        x7 = self.decoder_block1(x_center, x6)  # (bs, 384, T / 32, F' / 32)
+        x8 = self.decoder_block2(x7, x5)  # (bs, 384, T / 16, F' / 16)
+        x9 = self.decoder_block3(x8, x4)  # (bs, 256, T / 8, F' / 8)
+        x10 = self.decoder_block4(x9, x3)  # (bs, 128, T / 4, F' / 4)
+        x11 = self.decoder_block5(x10, x2)  # (bs, 64, T / 2, F' / 2)
+        x12 = self.decoder_block6(x11, x1)  # (bs, 32, T, F')
+        x = self.after_conv_block1(x12)  # (bs, 32, T, F')
+
+        x = self.after_conv2(x)
+        # (batch_size, target_sources_num * input_channles * self.K * subbands_num, T, F')
+
+        if self.subbands_num > 1:
+            x = self.subband.synthesis(x)
+            # (batch_size, target_sources_num * input_channles * self.K, T, F)
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 1024 -> 1025.
+
+        x = x[:, :, 0:origin_len, :]
+        # (batch_size, target_sources_num * input_channles * self.K, T, F)
+
+        audio_length = mixtures.shape[2]
+
+        separated_audio = self.feature_maps_to_wav(x, mag, sin_in, cos_in, audio_length)
+        # separated_audio: (batch_size, target_sources_num * input_channels, segments_num)
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/models/unet_subbandtime.py

@@ -0,0 +1,389 @@
+from typing import Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchlibrosa.stft import ISTFT, STFT, magphase
+
+from bytesep.models.pytorch_modules import Base, init_bn, init_layer
+from bytesep.models.subband_tools.pqmf import PQMF
+from bytesep.models.unet import ConvBlock, DecoderBlock, EncoderBlock
+
+
+class UNetSubbandTime(nn.Module, Base):
+    def __init__(self, input_channels: int, target_sources_num: int):
+        r"""Subband waveform UNet."""
+        super(UNetSubbandTime, self).__init__()
+
+        self.input_channels = input_channels
+        self.target_sources_num = target_sources_num
+
+        window_size = 512  # 2048 // 4
+        hop_size = 110  # 441 // 4
+        center = True
+        pad_mode = "reflect"
+        window = "hann"
+        activation = "leaky_relu"
+        momentum = 0.01
+
+        self.subbands_num = 4
+        self.K = 3  # outputs: |M|, cos∠M, sin∠M
+
+        self.downsample_ratio = 2 ** 6  # This number equals 2^{#encoder_blcoks}
+
+        self.pqmf = PQMF(
+            N=self.subbands_num,
+            M=64,
+            project_root='bytesep/models/subband_tools/filters',
+        )
+
+        self.stft = STFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.istft = ISTFT(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)
+
+        self.encoder_block1 = EncoderBlock(
+            in_channels=input_channels * self.subbands_num,
+            out_channels=32,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block2 = EncoderBlock(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block3 = EncoderBlock(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block4 = EncoderBlock(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block5 = EncoderBlock(
+            in_channels=256,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.encoder_block6 = EncoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            downsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.conv_block7 = ConvBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block1 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block2 = DecoderBlock(
+            in_channels=384,
+            out_channels=384,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block3 = DecoderBlock(
+            in_channels=384,
+            out_channels=256,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block4 = DecoderBlock(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+        self.decoder_block5 = DecoderBlock(
+            in_channels=128,
+            out_channels=64,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.decoder_block6 = DecoderBlock(
+            in_channels=64,
+            out_channels=32,
+            kernel_size=(3, 3),
+            upsample=(2, 2),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv_block1 = ConvBlock(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=(3, 3),
+            activation=activation,
+            momentum=momentum,
+        )
+
+        self.after_conv2 = nn.Conv2d(
+            in_channels=32,
+            out_channels=target_sources_num
+            * input_channels
+            * self.K
+            * self.subbands_num,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+            padding=(0, 0),
+            bias=True,
+        )
+
+        self.init_weights()
+
+    def init_weights(self):
+        r"""Initialize weights."""
+        init_bn(self.bn0)
+        init_layer(self.after_conv2)
+
+    def feature_maps_to_wav(
+        self,
+        input_tensor: torch.Tensor,
+        sp: torch.Tensor,
+        sin_in: torch.Tensor,
+        cos_in: torch.Tensor,
+        audio_length: int,
+    ) -> torch.Tensor:
+        r"""Convert feature maps to waveform.
+
+        Args:
+            input_tensor: (batch_size, target_sources_num * input_channels * self.K, time_steps, freq_bins)
+            sp: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            sin_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+            cos_in: (batch_size, target_sources_num * input_channels, time_steps, freq_bins)
+
+        Outputs:
+            waveform: (batch_size, target_sources_num * input_channels, segment_samples)
+        """
+        batch_size, _, time_steps, freq_bins = input_tensor.shape
+
+        x = input_tensor.reshape(
+            batch_size,
+            self.target_sources_num,
+            self.input_channels,
+            self.K,
+            time_steps,
+            freq_bins,
+        )
+        # x: (batch_size, target_sources_num, input_channles, K, time_steps, freq_bins)
+
+        mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])
+        _mask_real = torch.tanh(x[:, :, :, 1, :, :])
+        _mask_imag = torch.tanh(x[:, :, :, 2, :, :])
+        _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)
+        # mask_cos, mask_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        out_cos = (
+            cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin
+        )
+        out_sin = (
+            sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin
+        )
+        # out_cos: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+        # out_sin: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate |Y|.
+        out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag)
+        # out_mag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Calculate Y_{real} and Y_{imag} for ISTFT.
+        out_real = out_mag * out_cos
+        out_imag = out_mag * out_sin
+        # out_real, out_imag: (batch_size, target_sources_num, input_channles, time_steps, freq_bins)
+
+        # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.
+        shape = (
+            batch_size * self.target_sources_num * self.input_channels,
+            1,
+            time_steps,
+            freq_bins,
+        )
+        out_real = out_real.reshape(shape)
+        out_imag = out_imag.reshape(shape)
+
+        # ISTFT.
+        x = self.istft(out_real, out_imag, audio_length)
+        # (batch_size * target_sources_num * input_channels, segments_num)
+
+        # Reshape.
+        waveform = x.reshape(
+            batch_size, self.target_sources_num * self.input_channels, audio_length
+        )
+        # (batch_size, target_sources_num * input_channels, segments_num)
+
+        return waveform
+
+    def forward(self, input_dict: Dict) -> Dict:
+        """Forward data into the module.
+
+        Args:
+            input_dict: dict, e.g., {
+                waveform: (batch_size, input_channels, segment_samples),
+                ...,
+            }
+
+        Outputs:
+            output_dict: dict, e.g., {
+                'waveform': (batch_size, input_channels, segment_samples),
+                ...,
+            }
+        """
+        mixtures = input_dict['waveform']
+        # (batch_size, input_channels, segment_samples)
+
+        if self.subbands_num > 1:
+            subband_x = self.pqmf.analysis(mixtures)
+            # -- subband_x: (batch_size, input_channels * subbands_num, segment_samples)
+            # -- subband_x: (batch_size, subbands_num * input_channels, segment_samples)
+        else:
+            subband_x = mixtures
+
+        # from IPython import embed; embed(using=False); os._exit(0)
+        # import soundfile
+        # soundfile.write(file='_zz.wav', data=subband_x.data.cpu().numpy()[0, 2], samplerate=11025)
+
+        mag, cos_in, sin_in = self.wav_to_spectrogram_phase(subband_x)
+        # mag, cos_in, sin_in: (batch_size, input_channels * subbands_num, time_steps, freq_bins)
+
+        # Batch normalize on individual frequency bins.
+        x = mag.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        # (batch_size, input_channels * subbands_num, time_steps, freq_bins)
+
+        # Pad spectrogram to be evenly divided by downsample ratio.
+        origin_len = x.shape[2]
+        pad_len = (
+            int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio
+            - origin_len
+        )
+        x = F.pad(x, pad=(0, 0, 0, pad_len))
+        # x: (batch_size, input_channels * subbands_num, padded_time_steps, freq_bins)
+
+        # Let frequency bins be evenly divided by 2, e.g., 257 -> 256
+        x = x[..., 0 : x.shape[-1] - 1]  # (bs, input_channels, T, F)
+        # x: (batch_size, input_channels * subbands_num, padded_time_steps, freq_bins)
+
+        # UNet
+        (x1_pool, x1) = self.encoder_block1(x)  # x1_pool: (bs, 32, T / 2, F' / 2)
+        (x2_pool, x2) = self.encoder_block2(x1_pool)  # x2_pool: (bs, 64, T / 4, F' / 4)
+        (x3_pool, x3) = self.encoder_block3(
+            x2_pool
+        )  # x3_pool: (bs, 128, T / 8, F' / 8)
+        (x4_pool, x4) = self.encoder_block4(
+            x3_pool
+        )  # x4_pool: (bs, 256, T / 16, F' / 16)
+        (x5_pool, x5) = self.encoder_block5(
+            x4_pool
+        )  # x5_pool: (bs, 384, T / 32, F' / 32)
+        (x6_pool, x6) = self.encoder_block6(
+            x5_pool
+        )  # x6_pool: (bs, 384, T / 64, F' / 64)
+        x_center = self.conv_block7(x6_pool)  # (bs, 384, T / 64, F' / 64)
+        x7 = self.decoder_block1(x_center, x6)  # (bs, 384, T / 32, F' / 32)
+        x8 = self.decoder_block2(x7, x5)  # (bs, 384, T / 16, F' / 16)
+        x9 = self.decoder_block3(x8, x4)  # (bs, 256, T / 8, F' / 8)
+        x10 = self.decoder_block4(x9, x3)  # (bs, 128, T / 4, F' / 4)
+        x11 = self.decoder_block5(x10, x2)  # (bs, 64, T / 2, F' / 2)
+        x12 = self.decoder_block6(x11, x1)  # (bs, 32, T, F')
+        x = self.after_conv_block1(x12)  # (bs, 32, T, F')
+
+        x = self.after_conv2(x)
+        # (batch_size, subbands_num * target_sources_num * input_channles * self.K, T, F')
+
+        # Recover shape
+        x = F.pad(x, pad=(0, 1))  # Pad frequency, e.g., 256 -> 257.
+
+        x = x[:, :, 0:origin_len, :]
+        # (batch_size, subbands_num * target_sources_num * input_channles * self.K, T, F')
+
+        audio_length = subband_x.shape[2]
+
+        # Recover each subband spectrograms to subband waveforms. Then synthesis
+        # the subband waveforms to a waveform.
+        C1 = x.shape[1] // self.subbands_num
+        C2 = mag.shape[1] // self.subbands_num
+
+        separated_subband_audio = torch.cat(
+            [
+                self.feature_maps_to_wav(
+                    input_tensor=x[:, j * C1 : (j + 1) * C1, :, :],
+                    sp=mag[:, j * C2 : (j + 1) * C2, :, :],
+                    sin_in=sin_in[:, j * C2 : (j + 1) * C2, :, :],
+                    cos_in=cos_in[:, j * C2 : (j + 1) * C2, :, :],
+                    audio_length=audio_length,
+                )
+                for j in range(self.subbands_num)
+            ],
+            dim=1,
+        )
+        # （batch_size, subbands_num * target_sources_num * input_channles, segment_samples)
+
+        if self.subbands_num > 1:
+            separated_audio = self.pqmf.synthesis(separated_subband_audio)
+            # (batch_size, target_sources_num * input_channles, segment_samples)
+        else:
+            separated_audio = separated_subband_audio
+
+        output_dict = {'waveform': separated_audio}
+
+        return output_dict

bytesep/optimizers/lr_schedulers.py

@@ -0,0 +1,20 @@
+def get_lr_lambda(step, warm_up_steps: int, reduce_lr_steps: int):
+    r"""Get lr_lambda for LambdaLR. E.g.,
+
+    .. code-block: python
+        lr_lambda = lambda step: get_lr_lambda(step, warm_up_steps=1000, reduce_lr_steps=10000)
+
+        from torch.optim.lr_scheduler import LambdaLR
+        LambdaLR(optimizer, lr_lambda)
+
+    Args:
+        warm_up_steps: int, steps for warm up
+        reduce_lr_steps: int, reduce learning rate by 0.9 every #reduce_lr_steps steps
+
+    Returns:
+        learning rate: float
+    """
+    if step <= warm_up_steps:
+        return step / warm_up_steps
+    else:
+        return 0.9 ** (step // reduce_lr_steps)

bytesep/plot_results/musdb18.py

@@ -0,0 +1,198 @@
+import argparse
+import os
+import pickle
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def load_sdrs(workspace, task_name, filename, config, gpus, source_type):
+
+    stat_path = os.path.join(
+        workspace,
+        "statistics",
+        task_name,
+        filename,
+        "config={},gpus={}".format(config, gpus),
+        "statistics.pkl",
+    )
+
+    stat_dict = pickle.load(open(stat_path, 'rb'))
+
+    median_sdrs = [e['median_sdr_dict'][source_type] for e in stat_dict['test']]
+
+    return median_sdrs
+
+
+def plot_statistics(args):
+
+    # arguments & parameters
+    workspace = args.workspace
+    select = args.select
+    task_name = "musdb18"
+    filename = "train"
+
+    # paths
+    fig_path = os.path.join('results', task_name, "sdr_{}.pdf".format(select))
+    os.makedirs(os.path.dirname(fig_path), exist_ok=True)
+
+    linewidth = 1
+    lines = []
+    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
+
+    if select == '1a':
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='vocals-accompaniment,unet',
+            gpus=1,
+            source_type="vocals",
+        )
+        (line,) = ax.plot(sdrs, label='UNet,l1_wav', linewidth=linewidth)
+        lines.append(line)
+        ylim = 15
+
+    elif select == '1b':
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='accompaniment-vocals,unet',
+            gpus=1,
+            source_type="accompaniment",
+        )
+        (line,) = ax.plot(sdrs, label='UNet,l1_wav', linewidth=linewidth)
+        lines.append(line)
+        ylim = 20
+
+    if select == '1c':
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='vocals-accompaniment,unet',
+            gpus=1,
+            source_type="vocals",
+        )
+        (line,) = ax.plot(sdrs, label='UNet,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='vocals-accompaniment,resunet',
+            gpus=2,
+            source_type="vocals",
+        )
+        (line,) = ax.plot(sdrs, label='ResUNet_ISMIR2021,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='vocals-accompaniment,unet_subbandtime',
+            gpus=1,
+            source_type="vocals",
+        )
+        (line,) = ax.plot(sdrs, label='unet_subband,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='vocals-accompaniment,resunet_subbandtime',
+            gpus=1,
+            source_type="vocals",
+        )
+        (line,) = ax.plot(sdrs, label='resunet_subband,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        ylim = 15
+
+    elif select == '1d':
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='accompaniment-vocals,unet',
+            gpus=1,
+            source_type="accompaniment",
+        )
+        (line,) = ax.plot(sdrs, label='UNet,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='accompaniment-vocals,resunet',
+            gpus=2,
+            source_type="accompaniment",
+        )
+        (line,) = ax.plot(sdrs, label='ResUNet_ISMIR2021,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+        # sdrs = load_sdrs(
+        #     workspace,
+        #     task_name,
+        #     filename,
+        #     config='accompaniment-vocals,unet_subbandtime',
+        #     gpus=1,
+        #     source_type="accompaniment",
+        # )
+        # (line,) = ax.plot(sdrs, label='UNet_subbtandtime,l1_wav', linewidth=linewidth)
+        # lines.append(line)
+
+        sdrs = load_sdrs(
+            workspace,
+            task_name,
+            filename,
+            config='accompaniment-vocals,resunet_subbandtime',
+            gpus=1,
+            source_type="accompaniment",
+        )
+        (line,) = ax.plot(
+            sdrs, label='ResUNet_subbtandtime,l1_wav', linewidth=linewidth
+        )
+        lines.append(line)
+
+        ylim = 20
+
+    else:
+        raise Exception('Error!')
+
+    eval_every_iterations = 10000
+    total_ticks = 50
+    ticks_freq = 10
+
+    ax.set_ylim(0, ylim)
+    ax.set_xlim(0, total_ticks)
+    ax.xaxis.set_ticks(np.arange(0, total_ticks + 1, ticks_freq))
+    ax.xaxis.set_ticklabels(
+        np.arange(
+            0,
+            total_ticks * eval_every_iterations + 1,
+            ticks_freq * eval_every_iterations,
+        )
+    )
+    ax.yaxis.set_ticks(np.arange(ylim + 1))
+    ax.yaxis.set_ticklabels(np.arange(ylim + 1))
+    ax.grid(color='b', linestyle='solid', linewidth=0.3)
+    plt.legend(handles=lines, loc=4)
+
+    plt.savefig(fig_path)
+    print('Save figure to {}'.format(fig_path))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--workspace', type=str, required=True)
+    parser.add_argument('--select', type=str, required=True)
+
+    args = parser.parse_args()
+
+    plot_statistics(args)

bytesep/plot_results/plot_vctk-musdb18.py

@@ -0,0 +1,87 @@
+import os
+import sys
+import numpy as np
+import argparse
+import h5py
+import math
+import time
+import logging
+import pickle
+import matplotlib.pyplot as plt
+
+
+def load_sdrs(workspace, task_name, filename, config, gpus):
+
+    stat_path = os.path.join(
+        workspace,
+        "statistics",
+        task_name,
+        filename,
+        "config={},gpus={}".format(config, gpus),
+        "statistics.pkl",
+    )
+
+    stat_dict = pickle.load(open(stat_path, 'rb'))
+
+    median_sdrs = [e['sdr'] for e in stat_dict['test']]
+
+    return median_sdrs
+
+
+def plot_statistics(args):
+
+    # arguments & parameters
+    workspace = args.workspace
+    select = args.select
+    task_name = "vctk-musdb18"
+    filename = "train"
+
+    # paths
+    fig_path = os.path.join('results', task_name, "sdr_{}.pdf".format(select))
+    os.makedirs(os.path.dirname(fig_path), exist_ok=True)
+
+    linewidth = 1
+    lines = []
+    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
+    ylim = 30
+    expand = 1
+
+    if select == '1a':
+        sdrs = load_sdrs(workspace, task_name, filename, config='unet', gpus=1)
+        (line,) = ax.plot(sdrs, label='UNet,l1_wav', linewidth=linewidth)
+        lines.append(line)
+
+    else:
+        raise Exception('Error!')
+
+    eval_every_iterations = 10000
+    total_ticks = 50
+    ticks_freq = 10
+
+    ax.set_ylim(0, ylim)
+    ax.set_xlim(0, total_ticks)
+    ax.xaxis.set_ticks(np.arange(0, total_ticks + 1, ticks_freq))
+    ax.xaxis.set_ticklabels(
+        np.arange(
+            0,
+            total_ticks * eval_every_iterations + 1,
+            ticks_freq * eval_every_iterations,
+        )
+    )
+    ax.yaxis.set_ticks(np.arange(ylim + 1))
+    ax.yaxis.set_ticklabels(np.arange(ylim + 1))
+    ax.grid(color='b', linestyle='solid', linewidth=0.3)
+    plt.legend(handles=lines, loc=4)
+
+    plt.savefig(fig_path)
+    print('Save figure to {}'.format(fig_path))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--workspace', type=str, required=True)
+    parser.add_argument('--select', type=str, required=True)
+
+    args = parser.parse_args()
+
+    plot_statistics(args)