Noise-Aware Speech Separation with Contrastive Learning
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import sys
import hyperpyyaml

sys.path.insert(-1, os.path.abspath("../"))


# -- Project information -----------------------------------------------------

project = "SpeechBrain"
copyright = "2021, SpeechBrain"
author = "SpeechBrain"

# The full version, including alpha/beta/rc tags
release = "0.5.0"


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.intersphinx",
    "sphinx.ext.mathjax",
    "sphinx.ext.viewcode",
    "sphinx.ext.autosummary",
    "sphinx.ext.napoleon",
    "recommonmark",
]


# Napoleon settings
napoleon_google_docstring = False
napoleon_numpy_docstring = True
napoleon_include_init_with_doc = True
napoleon_include_private_with_doc = False
napoleon_include_special_with_doc = True
napoleon_use_admonition_for_examples = False
napoleon_use_admonition_for_notes = True
napoleon_use_admonition_for_references = False
napoleon_use_ivar = False
napoleon_use_param = True
napoleon_use_rtype = True

# Intersphinx mapping:
intersphinx_mapping = {
    "python": ("https://docs.python.org/", None),
    "numpy": ("http://docs.scipy.org/doc/numpy/", None),
    "torch": ("https://pytorch.org/docs/master/", None),
}

# AUTODOC:

autodoc_default_options = {}

# Autodoc mock extra dependencies:
autodoc_mock_imports = []

# Order of API items:
autodoc_member_order = "bysource"
autodoc_default_options = {"member-order": "bysource"}

# Don't show inherited docstrings:
autodoc_inherit_docstrings = False

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_apidoc_templates"]

# -- Better apidoc -----------------------------------------------------------


def run_apidoc(app):
    """Generage API documentation"""
    import better_apidoc

    better_apidoc.APP = app

    better_apidoc.main(
        [
            "better-apidoc",
            "-t",
            "_apidoc_templates",
            "--force",
            "--no-toc",
            "--separate",
            "-o",
            "API",
            os.path.dirname(hyperpyyaml.__file__),
        ]
    )
    better_apidoc.main(
        [
            "better-apidoc",
            "-t",
            "_apidoc_templates",
            "--force",
            "--no-toc",
            "--separate",
            "-o",
            "API",
            os.path.join("../", "speechbrain"),
        ]
    )


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = "sphinx_rtd_theme"
# See https://sphinx-rtd-theme.readthedocs.io/en/stable/configuring.html
# for rtd theme options
html_theme_options = {
    # Toc options
    "collapse_navigation": False,
    "sticky_navigation": True,
    "navigation_depth": 4,
    "includehidden": True,
}


# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

source_suffix = {
    ".rst": "restructuredtext",
    ".txt": "markdown",
    ".md": "markdown",
}


def setup(app):
    app.connect("builder-inited", run_apidoc)

docs/conf.py
TzuchengChang/NASS
import speechbrain as sb
import numpy as np
import torch
import torchaudio
import glob
import os
from speechbrain.dataio.batch import PaddedBatch
from tqdm import tqdm
import warnings
import pyloudnorm
import random

"""
The functions to implement Dynamic Mixing For SpeechSeparation

Authors
    * Samuele Cornell 2021
    * Cem Subakan 2021
"""


def build_spk_hashtable_librimix(hparams):
    """
    This function builds a dictionary of speaker-utterance pairs to be used in dynamic mixing
    """
    libri_utterances = glob.glob(
        os.path.join(hparams["base_folder_dm"], "**/*.wav"), recursive=True
    )

    spk_hashtable = {}

    # just for one file check if the sample rate is correct
    assert (
            torchaudio.info(libri_utterances[0]).sample_rate
            == hparams["sample_rate"]
    )
    for utt in tqdm(libri_utterances):

        path = os.path.normpath(utt)
        path_list = path.split(os.sep)
        spk_id = path_list[-3]

        # e.g. LibriSpeech/train-clean-360/441/128988/441-128988-0014.flac
        # id of speaker is 441 utterance is 128988-0014

        if spk_id not in spk_hashtable.keys():
            spk_hashtable[spk_id] = [utt]
        else:
            spk_hashtable[spk_id].append(utt)

    # calculate weights for each speaker ( len of list of utterances)
    spk_weights = [len(spk_hashtable[x]) for x in spk_hashtable.keys()]

    return spk_hashtable, spk_weights


def get_wham_noise_filenames(hparams):
    """This function lists the WHAM! noise files to be used in dynamic mixing"""

    if "Libri" in hparams["data_folder"]:
        # Data folder should point to Libri2Mix folder
        if hparams["sample_rate"] == 8000:
            noise_path = "wav8k/min/train-100/mynoise/"
        elif hparams["sample_rate"] == 16000:
            noise_path = "wav16k/min/train-100/mynoise/"
        else:
            raise ValueError("Unsupported Sampling Rate")
    else:
        if hparams["sample_rate"] == 8000:
            noise_path = "wav8k/min/tr/mynoise/"
        elif hparams["sample_rate"] == 16000:
            noise_path = "wav16k/min/tr/mynoise/"
        else:
            raise ValueError("Unsupported Sampling Rate")

    noise_files = glob.glob(
        os.path.join(hparams["data_folder"], noise_path, "*.wav")
    )
    return noise_files


def dynamic_mix_data_prep_librimix(hparams):
    """
    Dynamic mixing for LibriMix
    """

    # 1. Define datasets
    train_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["train_data"],
        replacements={"data_root": hparams["data_folder"]},
    )

    # we build an dictionary where keys are speakers id and entries are list
    # of utterances files of that speaker

    print("Building the speaker hashtable for dynamic mixing")
    spk_hashtable, spk_weights = build_spk_hashtable_librimix(hparams)

    spk_list = [x for x in spk_hashtable.keys()]
    spk_weights = [x / sum(spk_weights) for x in spk_weights]

    if hparams["use_wham_noise"]:
        noise_files = get_wham_noise_filenames(hparams)

    @sb.utils.data_pipeline.takes("mix_wav")
    @sb.utils.data_pipeline.provides(
        "mix_sig", "s1_sig", "s2_sig", "s3_sig", "noise_sig"
    )
    def audio_pipeline(
            mix_wav,
    ):  # this is dummy --> it means one epoch will be same as without dynamic mixing
        """
        This audio pipeline defines the compute graph for dynamic mixing
        """

        speakers = np.random.choice(
            spk_list, hparams["num_spks"], replace=False, p=spk_weights
        )

        if hparams["use_wham_noise"]:
            noise_file = np.random.choice(noise_files, 1, replace=False)

            noise, fs_read = torchaudio.load(noise_file[0])
            noise = noise.squeeze()

        # select two speakers randomly
        sources = []
        spk_files = [
            np.random.choice(spk_hashtable[spk], 1, False)[0]
            for spk in speakers
        ]

        if hparams["use_wham_noise"]:
            minlen = min(
                *[torchaudio.info(x).num_frames for x in spk_files],
                len(noise),
                hparams["training_signal_len"],
            )
        else:
            minlen = min(
                *[torchaudio.info(x).num_frames for x in spk_files],
                hparams["training_signal_len"],
            )

        meter = pyloudnorm.Meter(hparams["sample_rate"])

        MAX_AMP = 0.9
        MIN_LOUDNESS = -33
        MAX_LOUDNESS = -25

        def normalize(signal, is_noise=False):
            """
            This function normalizes the audio signals for loudness
            """
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                c_loudness = meter.integrated_loudness(signal)
                if is_noise:
                    target_loudness = random.uniform(
                        MIN_LOUDNESS - 5, MAX_LOUDNESS - 5
                        # MIN_LOUDNESS + 5, MAX_LOUDNESS + 5
                    )
                else:
                    target_loudness = random.uniform(MIN_LOUDNESS, MAX_LOUDNESS)
                signal = pyloudnorm.normalize.loudness(
                    signal, c_loudness, target_loudness
                )

                # check for clipping
                if np.max(np.abs(signal)) >= 1:
                    signal = signal * MAX_AMP / np.max(np.abs(signal))

            return torch.from_numpy(signal)

        for i, spk_file in enumerate(spk_files):
            # select random offset
            length = torchaudio.info(spk_file).num_frames
            start = 0
            stop = length
            if length > minlen:  # take a random window
                start = np.random.randint(0, length - minlen)
                stop = start + minlen

            tmp, fs_read = torchaudio.load(
                spk_file, frame_offset=start, num_frames=stop - start,
            )
            tmp = tmp[0].numpy()
            tmp = normalize(tmp)
            sources.append(tmp)

        sources = torch.stack(sources)
        mixture = torch.sum(sources, 0)
        if hparams["use_wham_noise"]:
            len_noise = len(noise)
            len_mix = len(mixture)
            min_len = min(len_noise, len_mix)
            noise = normalize(noise.numpy(), is_noise=True)

            start = 0
            stop = len_noise
            if len_noise > minlen:  # take a random window
                start = np.random.randint(0, len_noise - minlen)
                stop = start + minlen

            noise = noise[start:stop]
            # torchaudio.save(f"noise.wav", noise.unsqueeze(0), 8000)
            mixture = mixture[:min_len] + noise
            # torchaudio.save(f"mixture.wav", mixture.unsqueeze(0), 8000)

        # check for clipping
        max_amp_insig = mixture.abs().max().item()
        if max_amp_insig > MAX_AMP:
            weight = MAX_AMP / max_amp_insig
        else:
            weight = 1

        sources = weight * sources
        mixture = weight * mixture

        yield mixture
        for i in range(hparams["num_spks"]):
            yield sources[i]

        # If the number of speakers is 2, yield None for the 3rd speaker
        if hparams["num_spks"] == 2:
            yield None

        if hparams["use_wham_noise"]:
            noise = noise * weight
            yield noise
        else:
            yield None

    sb.dataio.dataset.add_dynamic_item([train_data], audio_pipeline)
    sb.dataio.dataset.set_output_keys(
        [train_data],
        ["id", "mix_sig", "s1_sig", "s2_sig", "s3_sig", "noise_sig"],
    )

    train_data = torch.utils.data.DataLoader(
        train_data,
        batch_size=hparams["dataloader_opts"]["batch_size"],
        num_workers=hparams["dataloader_opts"]["num_workers"],
        collate_fn=PaddedBatch,
        worker_init_fn=lambda x: np.random.seed(
            int.from_bytes(os.urandom(4), "little") + x
        ),
    )
    return train_data

recipes/LibriMix/separation/dynamic_mixing.py
TzuchengChang/NASS
"""
The functions to create the .csv files for LibriMix

Author
 * Cem Subakan 2020
"""

import os
import csv


def prepare_librimix(
        datapath,
        savepath,
        n_spks=2,
        skip_prep=False,
        librimix_addnoise=False,
        fs=8000,
):
    """

    Prepare .csv files for librimix

    Arguments:
    ----------
        datapath (str) : path for the wsj0-mix dataset.
        savepath (str) : path where we save the csv file.
        n_spks (int): number of speakers
        skip_prep (bool): If True, skip data preparation
        librimix_addnoise: If True, add whamnoise to librimix datasets
    """

    if skip_prep:
        return

    if "Libri" in datapath:
        # Libri 2/3Mix datasets
        if n_spks == 2:
            assert (
                    "Libri2Mix" in datapath
            ), "Inconsistent number of speakers and datapath"
            create_libri2mix_csv(datapath, savepath, addnoise=librimix_addnoise)
        elif n_spks == 3:
            assert (
                    "Libri3Mix" in datapath
            ), "Inconsistent number of speakers and datapath"
            create_libri3mix_csv(datapath, savepath, addnoise=librimix_addnoise)
        else:
            raise ValueError("Unsupported Number of Speakers")
    else:
        print("Creating a csv file for a custom dataset")
        create_custom_dataset(datapath, savepath)


def create_libri2mix_csv(
        datapath,
        savepath,
        addnoise=False,
        version="wav8k/min/",
        set_types=["train-100", "dev", "test"],
):
    """
    This functions creates the .csv file for the libri2mix dataset
    """

    for set_type in set_types:
        if addnoise:
            mix_path = os.path.join(datapath, version, set_type, "mix_both/")
        else:
            mix_path = os.path.join(datapath, version, set_type, "mix_clean/")

        s1_path = os.path.join(datapath, version, set_type, "s1/")
        s2_path = os.path.join(datapath, version, set_type, "s2/")
        noise_path = os.path.join(datapath, version, set_type, "noise/")

        files = os.listdir(mix_path)

        mix_fl_paths = [mix_path + fl for fl in files]
        s1_fl_paths = [s1_path + fl for fl in files]
        s2_fl_paths = [s2_path + fl for fl in files]
        noise_fl_paths = [noise_path + fl for fl in files]

        csv_columns = [
            "ID",
            "duration",
            "mix_wav",
            "mix_wav_format",
            "mix_wav_opts",
            "s1_wav",
            "s1_wav_format",
            "s1_wav_opts",
            "s2_wav",
            "s2_wav_format",
            "s2_wav_opts",
            "noise_wav",
            "noise_wav_format",
            "noise_wav_opts",
        ]

        with open(savepath + "/libri2mix_" + set_type + ".csv", "w") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
            writer.writeheader()
            for i, (mix_path, s1_path, s2_path, noise_path) in enumerate(
                    zip(mix_fl_paths, s1_fl_paths, s2_fl_paths, noise_fl_paths)
            ):
                row = {
                    "ID": i,
                    "duration": 1.0,
                    "mix_wav": mix_path,
                    "mix_wav_format": "wav",
                    "mix_wav_opts": None,
                    "s1_wav": s1_path,
                    "s1_wav_format": "wav",
                    "s1_wav_opts": None,
                    "s2_wav": s2_path,
                    "s2_wav_format": "wav",
                    "s2_wav_opts": None,
                    "noise_wav": noise_path,
                    "noise_wav_format": "wav",
                    "noise_wav_opts": None,
                }
                writer.writerow(row)


def create_libri3mix_csv(
        datapath,
        savepath,
        addnoise=False,
        version="wav8k/min/",
        set_types=["train-100", "dev", "test"],
):
    """
    This functions creates the .csv file for the libri3mix dataset
    """

    for set_type in set_types:
        if addnoise:
            mix_path = os.path.join(datapath, version, set_type, "mix_both/")
        else:
            mix_path = os.path.join(datapath, version, set_type, "mix_clean/")

        s1_path = os.path.join(datapath, version, set_type, "s1/")
        s2_path = os.path.join(datapath, version, set_type, "s2/")
        s3_path = os.path.join(datapath, version, set_type, "s3/")
        noise_path = os.path.join(datapath, version, set_type, "noise/")

        files = os.listdir(mix_path)

        mix_fl_paths = [mix_path + fl for fl in files]
        s1_fl_paths = [s1_path + fl for fl in files]
        s2_fl_paths = [s2_path + fl for fl in files]
        s3_fl_paths = [s3_path + fl for fl in files]
        noise_fl_paths = [noise_path + fl for fl in files]

        csv_columns = [
            "ID",
            "duration",
            "mix_wav",
            "mix_wav_format",
            "mix_wav_opts",
            "s1_wav",
            "s1_wav_format",
            "s1_wav_opts",
            "s2_wav",
            "s2_wav_format",
            "s2_wav_opts",
            "s3_wav",
            "s3_wav_format",
            "s3_wav_opts",
            "noise_wav",
            "noise_wav_format",
            "noise_wav_opts",
        ]

        with open(savepath + "/libri3mix_" + set_type + ".csv", "w") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
            writer.writeheader()
            for (
                    i,
                    (mix_path, s1_path, s2_path, s3_path, noise_path),
            ) in enumerate(
                zip(
                    mix_fl_paths,
                    s1_fl_paths,
                    s2_fl_paths,
                    s3_fl_paths,
                    noise_fl_paths,
                )
            ):
                row = {
                    "ID": i,
                    "duration": 1.0,
                    "mix_wav": mix_path,
                    "mix_wav_format": "wav",
                    "mix_wav_opts": None,
                    "s1_wav": s1_path,
                    "s1_wav_format": "wav",
                    "s1_wav_opts": None,
                    "s2_wav": s2_path,
                    "s2_wav_format": "wav",
                    "s2_wav_opts": None,
                    "s3_wav": s3_path,
                    "s3_wav_format": "wav",
                    "s3_wav_opts": None,
                    "noise_wav": noise_path,
                    "noise_wav_format": "wav",
                    "noise_wav_opts": None,
                }
                writer.writerow(row)


def create_custom_dataset(
        datapath,
        savepath,
        dataset_name="custom",
        set_types=["train", "valid", "test"],
        folder_names={
            "source1": "source1/",
            "source2": "source2/",
            "source3": "source3/",
            "mixture": "mixture/",
            "noise": "noise/",
        },
):
    """
    This function creates the csv file for a custom source separation dataset
    """

    for set_type in set_types:
        mix_path = os.path.join(datapath, set_type, folder_names["mixture"])
        s1_path = os.path.join(datapath, set_type, folder_names["source1"])
        s2_path = os.path.join(datapath, set_type, folder_names["source2"])
        s3_path = os.path.join(datapath, set_type, folder_names["source3"])
        noise_path = os.path.join(datapath, set_type, folder_names["noise"])

        files = os.listdir(mix_path)

        mix_fl_paths = [os.path.join(mix_path, fl) for fl in files]
        s1_fl_paths = [os.path.join(s1_path, fl) for fl in files]
        s2_fl_paths = [os.path.join(s2_path, fl) for fl in files]
        s3_fl_paths = [s3_path + fl for fl in files]
        noise_fl_paths = [noise_path + fl for fl in files]

        csv_columns = [
            "ID",
            "duration",
            "mix_wav",
            "mix_wav_format",
            "mix_wav_opts",
            "s1_wav",
            "s1_wav_format",
            "s1_wav_opts",
            "s2_wav",
            "s2_wav_format",
            "s2_wav_opts",
            "s3_wav",
            "s3_wav_format",
            "s3_wav_opts",
            "noise_wav",
            "noise_wav_format",
            "noise_wav_opts",
        ]

        with open(
                os.path.join(savepath, dataset_name + "_" + set_type + ".csv"), "w"
        ) as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
            writer.writeheader()
            for i, (mix_path, s1_path, s2_path, s3_path, noise_path) in enumerate(
                    zip(mix_fl_paths, s1_fl_paths, s2_fl_paths,
                        s3_fl_paths,
                        noise_fl_paths)
            ):
                row = {
                    "ID": i,
                    "duration": 1.0,
                    "mix_wav": mix_path,
                    "mix_wav_format": "wav",
                    "mix_wav_opts": None,
                    "s1_wav": s1_path,
                    "s1_wav_format": "wav",
                    "s1_wav_opts": None,
                    "s2_wav": s2_path,
                    "s2_wav_format": "wav",
                    "s2_wav_opts": None,
                    "s3_wav": s3_path,
                    "s3_wav_format": "wav",
                    "s3_wav_opts": None,
                    "noise_wav": noise_path,
                    "noise_wav_format": "wav",
                    "noise_wav_opts": None,
                }
                writer.writerow(row)

recipes/LibriMix/separation/prepare_data.py
TzuchengChang/NASS
"""
This script allows to resample a folder which contains audio files.
The files are parsed recursively. An exact copy of the folder is created,
with same structure but contained resampled audio files.
Resampling is performed by using sox through torchaudio.
Author
------
Samuele Cornell, 2020
"""

import os
import argparse
from pathlib import Path
import tqdm
import torchaudio
import glob

# from oct2py import octave
from scipy import signal
import numpy as np
import torch

parser = argparse.ArgumentParser(
    "utility for resampling all audio files in a folder recursively"
    "It --input_folder to --output_folder and "
    "resamples all audio files with specified format to --fs."
)
parser.add_argument("--input_folder", type=str, required=True)
parser.add_argument("--output_folder", type=str, required=True)
parser.add_argument(
    "--fs", type=str, default=8000, help="this is the target sampling frequency"
)
parser.add_argument("--regex", type=str, default="**/*.flac")


def resample_folder(input_folder, output_folder, fs, regex):
    """Resamples the wav files within an input folder.

    Arguments
    ---------
    input_folder : path
        Path of the folder to resample.
    output_folder : path
        Path of the output folder with the resampled data.
    fs : int
        Target sampling frequency.
    regex: str
        Regular expression for search.
    """
    # filedir = os.path.dirname(os.path.realpath(__file__))
    # octave.addpath(filedir)
    # add the matlab functions to octave dir here

    files = glob.glob(os.path.join(input_folder, regex), recursive=True)
    for f in tqdm.tqdm(files):
        audio, fs_read = torchaudio.load(f)
        audio = audio[0].numpy()
        audio = signal.resample_poly(audio, fs, fs_read)

        # tmp = octave.activlev(audio.tolist(), fs, "n")
        # audio, _ = tmp[:-1].squeeze(), tmp[-1]

        peak = np.max(np.abs(audio))
        audio = audio / peak
        audio = torch.from_numpy(audio).float()

        relative_path = os.path.join(
            Path(f).relative_to(Path(input_folder)).parent,
            Path(f).relative_to(Path(input_folder)).stem
            + "_peak_{}.wav".format(peak),
        )

        os.makedirs(
            Path(
                os.path.join(
                    output_folder, Path(f).relative_to(Path(input_folder))
                )
            ).parent,
            exist_ok=True,
        )

        torchaudio.save(
            os.path.join(output_folder, relative_path),
            audio.reshape(1, -1),
            fs,
        )


if __name__ == "__main__":
    args = parser.parse_args()
    resample_folder(
        args.input_folder, args.output_folder, int(args.fs), args.regex
    )

recipes/LibriMix/separation/preprocess_dynamic_mixing.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""Recipe for training a neural speech separation system on Libri2/3Mix datasets.
The system employs an encoder, a decoder, and a masking network.

To run this recipe, do the following:
> python train.py hparams/sepformer-libri2mix.yaml
> python train.py hparams/sepformer-libri3mix.yaml


The experiment file is flexible enough to support different neural
networks. By properly changing the parameter files, you can try
different architectures. The script supports both libri2mix and
libri3mix.


Authors
 * Cem Subakan 2020
 * Mirco Ravanelli 2020
 * Samuele Cornell 2020
 * Mirko Bronzi 2020
 * Jianyuan Zhong 2020
"""

import os
import sys
import torch
import torch.nn.functional as F
import torchaudio
import speechbrain as sb
import speechbrain.nnet.schedulers as schedulers
from speechbrain.utils.distributed import run_on_main
from torch.cuda.amp import autocast
from hyperpyyaml import load_hyperpyyaml
import numpy as np
from tqdm import tqdm
import csv
import logging


# Define training procedure
class Separation(sb.Brain):
    def compute_forward(self, mix, targets, stage, noise=None):
        """Forward computations from the mixture to the separated signals."""

        # Unpack lists and put tensors in the right device
        mix, mix_lens = mix
        mix, mix_lens = mix.to(self.device), mix_lens.to(self.device)

        # Convert targets to tensor
        targets = torch.cat(
            [targets[i][0].unsqueeze(-1) for i in
             range(self.hparams.num_spks + 1 if self.hparams.use_noise_ground else self.hparams.num_spks)],
            dim=-1,
        ).to(self.device)

        # Add speech distortions
        if stage == sb.Stage.TRAIN:
            with torch.no_grad():
                if self.hparams.use_speedperturb or self.hparams.use_rand_shift:
                    mix, targets = self.add_speed_perturb(targets, mix_lens)

                    if self.hparams.use_wham_noise:
                        noise = noise.to(self.device)
                        len_noise = noise.shape[1]
                        len_mix = mix.shape[1]
                        min_len = min(len_noise, len_mix)

                        # add the noise
                        if self.hparams.use_noise_ground:
                            mix = mix[:, :min_len]
                        else:
                            mix = mix[:, :min_len] + noise[:, :min_len]

                        # fix the length of targets also
                        targets = targets[:, :min_len, :]

                if self.hparams.use_wavedrop:
                    mix = self.hparams.wavedrop(mix, mix_lens)

                if self.hparams.limit_training_signal_len:
                    mix, targets = self.cut_signals(mix, targets)

        # Separation
        mix_w = self.hparams.Encoder(mix)
        est_mask = self.hparams.MaskNet(mix_w)

        mix_w = torch.stack([mix_w] * (
                self.hparams.num_spks + 1)) if self.hparams.use_noise_ground and not self.hparams.use_noise_ground_only else torch.stack(
            [mix_w] * self.hparams.num_spks)

        sep_h = mix_w * est_mask

        if self.hparams.use_nceloss:
            targets_f = [self.hparams.Encoder(item.squeeze(-1)).unsqueeze(0) for item in
                         targets.chunk(self.hparams.num_spks + 1, -1)]
            estimates_f = [item for item in sep_h.chunk(self.hparams.num_spks + 1, 0)]
            _, ids = self.hparams.PatchSampleNet(estimates_f[:1], self.hparams.num_patches)
            ids *= (self.hparams.num_spks + 1)
            spk_ef_pool, _ = self.hparams.PatchSampleNet(estimates_f, self.hparams.num_patches, ids)
            spk_tf_pool, _ = self.hparams.PatchSampleNet(targets_f, self.hparams.num_patches, ids)

        # Decoding
        est_source = torch.cat(
            [
                self.hparams.Decoder(sep_h[i]).unsqueeze(-1)
                for i in range(
                self.hparams.num_spks + 1 if self.hparams.use_noise_ground and not self.hparams.use_noise_ground_only else self.hparams.num_spks)
            ],
            dim=-1,
        )

        # T changed after conv1d in encoder, fix it here
        T_origin = mix.size(1)
        T_est = est_source.size(1)
        if T_origin > T_est:
            est_source = F.pad(est_source, (0, 0, 0, T_origin - T_est))
        else:
            est_source = est_source[:, :T_origin, :]

        if self.hparams.use_nceloss:
            return est_source, targets, spk_ef_pool, spk_tf_pool
        else:
            return est_source, targets

    def compute_objectives(self, predictions, targets):
        """Computes the si-snr loss"""
        return self.hparams.loss(targets, predictions)

    def fit_batch(self, batch):
        """Trains one batch"""
        # Unpacking batch list
        mixture = batch.mix_sig
        targets = [batch.s1_sig, batch.s2_sig]
        if self.hparams.use_wham_noise:
            noise = batch.noise_sig[0]
        else:
            noise = None

        if self.hparams.num_spks == 3:
            targets.append(batch.s3_sig)

        if self.hparams.use_noise_ground:
            targets.append(batch.noise_sig)

        if self.auto_mix_prec:
            with autocast():
                if self.hparams.use_nceloss:
                    predictions, targets, spk_ef_pool, spk_tf_pool = self.compute_forward(mixture, targets,
                                                                                          sb.Stage.TRAIN,
                                                                                          noise)
                else:
                    predictions, targets = self.compute_forward(mixture, targets, sb.Stage.TRAIN, noise)

                if self.hparams.use_noise_ground and self.hparams.use_noise_ground_only:
                    loss = self.compute_objectives(predictions, targets[:, :, :-1])
                else:
                    loss, perms = self.compute_objectives(predictions, targets)

                if self.hparams.use_myloss:
                    if self.hparams.use_nceloss:
                        spk_ef_pool = [spk_ef_pool[i] for i in perms[0]]
                        total_nce_loss = 0.0
                        for f_q, f_p, f_n in zip(spk_ef_pool[:-1], spk_tf_pool[:-1],
                                                 spk_ef_pool[-1:] * self.hparams.num_spks):
                            temp = self.hparams.nceloss(f_q, f_p, f_n) * self.hparams.lambda_NCE
                            total_nce_loss += temp.mean()
                        loss = loss + total_nce_loss / self.hparams.num_spks

                # hard threshold the easy dataitems
                if self.hparams.threshold_byloss:
                    th = self.hparams.threshold
                    loss_to_keep = loss[loss > th]
                    if loss_to_keep.nelement() > 0:
                        loss = loss_to_keep.mean()
                else:
                    loss = loss.mean()

            if (
                    loss < self.hparams.loss_upper_lim and loss.nelement() > 0
            ):  # the fix for computational problems
                self.scaler.scale(loss).backward()
                if self.hparams.clip_grad_norm >= 0:
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(
                        self.modules.parameters(), self.hparams.clip_grad_norm,
                    )
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                self.nonfinite_count += 1
                logger.info(
                    "infinite loss or empty loss! it happened {} times so far - skipping this batch".format(
                        self.nonfinite_count
                    )
                )
                loss.data = torch.tensor(0.0).to(self.device)
        else:
            if self.hparams.use_nceloss:
                predictions, targets, spk_ef_pool, spk_tf_pool = self.compute_forward(mixture, targets,
                                                                                      sb.Stage.TRAIN,
                                                                                      noise)
            else:
                predictions, targets = self.compute_forward(mixture, targets, sb.Stage.TRAIN, noise)

            loss, perms = self.compute_objectives(predictions, targets)

            if self.hparams.threshold_byloss:
                th = self.hparams.threshold
                loss_to_keep = loss[loss > th]
                if loss_to_keep.nelement() > 0:
                    loss = loss_to_keep.mean()
            else:
                loss = loss.mean()

            if (
                    loss < self.hparams.loss_upper_lim and loss.nelement() > 0
            ):  # the fix for computational problems
                loss.backward()
                if self.hparams.clip_grad_norm >= 0:
                    torch.nn.utils.clip_grad_norm_(
                        self.modules.parameters(), self.hparams.clip_grad_norm
                    )
                self.optimizer.step()
            else:
                self.nonfinite_count += 1
                logger.info(
                    "infinite loss or empty loss! it happened {} times so far - skipping this batch".format(
                        self.nonfinite_count
                    )
                )
                loss.data = torch.tensor(0.0).to(self.device)
        self.optimizer.zero_grad()

        return loss.detach().cpu()

    def evaluate_batch(self, batch, stage):
        """Computations needed for validation/test batches"""
        snt_id = batch.id
        mixture = batch.mix_sig
        targets = [batch.s1_sig, batch.s2_sig]
        if self.hparams.num_spks == 3:
            targets.append(batch.s3_sig)

        if self.hparams.use_noise_ground:
            targets.append(batch.noise_sig)

        with torch.no_grad():
            if self.hparams.use_nceloss:
                predictions, targets, spk_ef_pool, spk_tf_pool = self.compute_forward(mixture, targets,
                                                                                      stage)
            else:
                predictions, targets = self.compute_forward(mixture, targets, stage)

            if self.hparams.use_noise_ground and self.hparams.use_noise_ground_only:
                loss = self.compute_objectives(predictions, targets[:, :, :-1])
            elif self.hparams.use_noise_ground and self.hparams.use_noise_notest:
                _, perms = self.compute_objectives(predictions, targets)
                predictions[:, :, list(perms[0])] = predictions.clone()
                loss, _ = self.compute_objectives(predictions[:, :, :-1], targets[:, :, :-1])
            else:
                loss, _ = self.compute_objectives(predictions, targets)

        # Manage audio file saving
        if stage == sb.Stage.TEST and self.hparams.save_audio:
            if hasattr(self.hparams, "n_audio_to_save"):
                if self.hparams.n_audio_to_save > 0:
                    self.save_audio(snt_id[0], mixture, targets, predictions)
                    self.hparams.n_audio_to_save += -1
            else:
                self.save_audio(snt_id[0], mixture, targets, predictions)

        return loss.mean().detach()

    def on_stage_end(self, stage, stage_loss, epoch):
        """Gets called at the end of a epoch."""
        # Compute/store important stats
        stage_stats = {"si-snr": stage_loss}
        if stage == sb.Stage.TRAIN:
            self.train_stats = stage_stats

        # Perform end-of-iteration things, like annealing, logging, etc.
        if stage == sb.Stage.VALID:

            # Learning rate annealing
            if isinstance(
                    self.hparams.lr_scheduler, schedulers.ReduceLROnPlateau
            ):
                current_lr, next_lr = self.hparams.lr_scheduler(
                    [self.optimizer], epoch, stage_loss
                )
                schedulers.update_learning_rate(self.optimizer, next_lr)
            else:
                # if we do not use the reducelronplateau, we do not change the lr
                current_lr = self.hparams.optimizer.optim.param_groups[0]["lr"]

            self.hparams.train_logger.log_stats(
                stats_meta={"epoch": epoch, "lr": current_lr},
                train_stats=self.train_stats,
                valid_stats=stage_stats,
            )
            self.checkpointer.save_and_keep_only(
                meta={"si-snr": stage_stats["si-snr"]}, min_keys=["si-snr"],
            )
        elif stage == sb.Stage.TEST:
            self.hparams.train_logger.log_stats(
                stats_meta={"Epoch loaded": self.hparams.epoch_counter.current},
                test_stats=stage_stats,
            )

    def add_speed_perturb(self, targets, targ_lens):
        """Adds speed perturbation and random_shift to the input signals"""

        min_len = -1

        if self.hparams.use_speedperturb:
            # Performing speed change (independently on each source)
            new_targets = []
            recombine = True

            for i in range(self.hparams.num_spks):
                new_target = self.hparams.speedperturb(
                    targets[:, :, i], targ_lens
                )
                new_targets.append(new_target)
                if i == 0:
                    min_len = new_target.shape[-1]
                else:
                    if new_target.shape[-1] < min_len:
                        min_len = new_target.shape[-1]

            if self.hparams.use_noise_ground:
                copy = targets[:, :, -1]
                new_targets.append(copy)
                if copy.shape[-1] < min_len:
                    min_len = copy.shape[-1]

            if self.hparams.use_rand_shift:
                # Performing random_shift (independently on each source)
                recombine = True
                for i in range(targets.shape[-1]):
                    rand_shift = torch.randint(
                        self.hparams.min_shift, self.hparams.max_shift, (1,)
                    )
                    new_targets[i] = new_targets[i].to(self.device)
                    new_targets[i] = torch.roll(
                        new_targets[i], shifts=(rand_shift[0],), dims=1
                    )

            # Re-combination
            if recombine:
                if self.hparams.use_speedperturb:
                    targets = torch.zeros(
                        targets.shape[0],
                        min_len,
                        targets.shape[-1],
                        device=targets.device,
                        dtype=torch.float,
                    )
                for i, new_target in enumerate(new_targets):
                    targets[:, :, i] = new_targets[i][:, 0:min_len]

        mix = targets.sum(-1)
        return mix, targets

    def cut_signals(self, mixture, targets):
        """This function selects a random segment of a given length within the mixture.
        The corresponding targets are selected accordingly"""

        min_len = min(mixture.shape[1], targets.shape[1], self.hparams.training_signal_len)

        randstart = torch.randint(
            0,
            1 + max(0, mixture.shape[1] - min_len),
            (1,),
        ).item()
        targets = targets[
                  :, randstart: randstart + min_len, :
                  ]
        mixture = mixture[
                  :, randstart: randstart + min_len
                  ]
        return mixture, targets

    def reset_layer_recursively(self, layer):
        """Reinitializes the parameters of the neural networks"""
        if hasattr(layer, "reset_parameters"):
            layer.reset_parameters()
        for child_layer in layer.modules():
            if layer != child_layer:
                self.reset_layer_recursively(child_layer)

    def save_results(self, test_data):
        """This script computes the SDR and SI-SNR metrics and saves
        them into a csv file"""

        # This package is required for SDR computation
        from mir_eval.separation import bss_eval_sources

        # Create folders where to store audio
        save_file = os.path.join(self.hparams.output_folder, "test_results.csv")

        # Variable init
        all_sdrs = []
        all_sdrs_i = []
        all_sisnrs = []
        all_sisnrs_i = []
        csv_columns = ["snt_id", "sdr", "sdr_i", "si-snr", "si-snr_i"]

        test_loader = sb.dataio.dataloader.make_dataloader(
            test_data, **self.hparams.dataloader_opts
        )

        with open(save_file, "w") as results_csv:
            writer = csv.DictWriter(results_csv, fieldnames=csv_columns)
            writer.writeheader()

            # Loop over all test sentence
            with tqdm(test_loader, dynamic_ncols=True) as t:
                for i, batch in enumerate(t):

                    # Apply Separation
                    mixture, mix_len = batch.mix_sig
                    snt_id = batch.id
                    targets = [batch.s1_sig, batch.s2_sig]
                    if self.hparams.num_spks == 3:
                        targets.append(batch.s3_sig)

                    if self.hparams.use_noise_ground:
                        targets.append(batch.noise_sig)

                    with torch.no_grad():
                        if self.hparams.use_nceloss:
                            predictions, targets, spk_ef_pool, spk_tf_pool = self.compute_forward(
                                batch.mix_sig, targets, sb.Stage.TEST
                            )
                        else:
                            predictions, targets = self.compute_forward(batch.mix_sig, targets, sb.Stage.TEST)

                    if self.hparams.use_noise_ground and self.hparams.use_noise_ground_only:
                        targets = targets[:, :, :-1]

                    if self.hparams.use_noise_ground and self.hparams.use_noise_notest:
                        _, perms = self.compute_objectives(predictions, targets)
                        predictions[:, :, list(perms[0])] = predictions.clone()
                        predictions = predictions[:, :, :-1]
                        targets = targets[:, :, :-1]

                    # Compute SI-SNR
                    sisnr, _ = self.compute_objectives(predictions, targets)

                    num_spks = self.hparams.num_spks + 1 if self.hparams.use_noise_ground and not self.hparams.use_noise_ground_only and not self.hparams.use_noise_notest else self.hparams.num_spks

                    # Compute SI-SNR improvement
                    mixture_signal = torch.stack(
                        [mixture] * num_spks, dim=-1
                    )

                    mixture_signal = mixture_signal.to(targets.device)
                    sisnr_baseline, _ = self.compute_objectives(
                        mixture_signal, targets
                    )
                    sisnr_i = sisnr - sisnr_baseline

                    # Compute SDR
                    sdr, _, _, _ = bss_eval_sources(
                        targets[0].t().cpu().numpy(),
                        predictions[0].t().detach().cpu().numpy(),
                    )

                    sdr_baseline, _, _, _ = bss_eval_sources(
                        targets[0].t().cpu().numpy(),
                        mixture_signal[0].t().detach().cpu().numpy(),
                    )

                    sdr_i = sdr.mean() - sdr_baseline.mean()

                    # Saving on a csv file
                    row = {
                        "snt_id": snt_id[0],
                        "sdr": sdr.mean(),
                        "sdr_i": sdr_i,
                        "si-snr": -sisnr.item(),
                        "si-snr_i": -sisnr_i.item(),
                    }
                    writer.writerow(row)

                    # Metric Accumulation
                    all_sdrs.append(sdr.mean())
                    all_sdrs_i.append(sdr_i.mean())
                    all_sisnrs.append(-sisnr.item())
                    all_sisnrs_i.append(-sisnr_i.item())

                    # break

                row = {
                    "snt_id": "avg",
                    "sdr": np.array(all_sdrs).mean(),
                    "sdr_i": np.array(all_sdrs_i).mean(),
                    "si-snr": np.array(all_sisnrs).mean(),
                    "si-snr_i": np.array(all_sisnrs_i).mean(),
                }
                writer.writerow(row)

        logger.info("Mean SISNR is {}".format(np.array(all_sisnrs).mean()))
        logger.info("Mean SISNRi is {}".format(np.array(all_sisnrs_i).mean()))
        logger.info("Mean SDR is {}".format(np.array(all_sdrs).mean()))
        logger.info("Mean SDRi is {}".format(np.array(all_sdrs_i).mean()))

    def save_audio(self, snt_id, mixture, targets, predictions):
        """saves the test audio (mixture, targets, and estimated sources) on disk"""

        # Create outout folder
        save_path = os.path.join(self.hparams.save_folder, "audio_results")
        if not os.path.exists(save_path):
            os.mkdir(save_path)

        for ns in range(
                self.hparams.num_spks + 1 if self.hparams.use_noise_ground and not self.hparams.use_noise_ground_only else self.hparams.num_spks):
            # Estimated source
            signal = predictions[0, :, ns]
            signal = signal / signal.abs().max()
            save_file = os.path.join(
                save_path, "item{}_source{}hat.wav".format(snt_id, ns + 1)
            )
            torchaudio.save(
                save_file, signal.unsqueeze(0).cpu(), self.hparams.sample_rate
            )

            # Original source
            signal = targets[0, :, ns]
            signal = signal / signal.abs().max()
            save_file = os.path.join(
                save_path, "item{}_source{}.wav".format(snt_id, ns + 1)
            )
            torchaudio.save(
                save_file, signal.unsqueeze(0).cpu(), self.hparams.sample_rate
            )

        # Mixture
        signal = mixture[0][0, :]
        signal = signal / signal.abs().max()
        save_file = os.path.join(save_path, "item{}_mix.wav".format(snt_id))
        torchaudio.save(
            save_file, signal.unsqueeze(0).cpu(), self.hparams.sample_rate
        )


def dataio_prep(hparams):
    """Creates data processing pipeline"""

    # 1. Define datasets
    train_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["train_data"],
        replacements={"data_root": hparams["data_folder"]},
    )

    valid_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["valid_data"],
        replacements={"data_root": hparams["data_folder"]},
    )

    test_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["test_data"],
        replacements={"data_root": hparams["data_folder"]},
    )

    datasets = [train_data, valid_data, test_data]

    # 2. Provide audio pipelines

    @sb.utils.data_pipeline.takes("mix_wav")
    @sb.utils.data_pipeline.provides("mix_sig")
    def audio_pipeline_mix(mix_wav):
        mix_sig = sb.dataio.dataio.read_audio(mix_wav)
        return mix_sig

    @sb.utils.data_pipeline.takes("s1_wav")
    @sb.utils.data_pipeline.provides("s1_sig")
    def audio_pipeline_s1(s1_wav):
        s1_sig = sb.dataio.dataio.read_audio(s1_wav)
        return s1_sig

    @sb.utils.data_pipeline.takes("s2_wav")
    @sb.utils.data_pipeline.provides("s2_sig")
    def audio_pipeline_s2(s2_wav):
        s2_sig = sb.dataio.dataio.read_audio(s2_wav)
        return s2_sig

    if hparams["num_spks"] == 3:
        @sb.utils.data_pipeline.takes("s3_wav")
        @sb.utils.data_pipeline.provides("s3_sig")
        def audio_pipeline_s3(s3_wav):
            s3_sig = sb.dataio.dataio.read_audio(s3_wav)
            return s3_sig

    if hparams["use_wham_noise"]:
        @sb.utils.data_pipeline.takes("noise_wav")
        @sb.utils.data_pipeline.provides("noise_sig")
        def audio_pipeline_noise(noise_wav):
            noise_sig = sb.dataio.dataio.read_audio(noise_wav)
            return noise_sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline_mix)
    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline_s1)
    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline_s2)
    if hparams["num_spks"] == 3:
        sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline_s3)

    if hparams["use_wham_noise"]:
        print("Using the WHAM! noise in the data pipeline")
        sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline_noise)

    if (hparams["num_spks"] == 2) and hparams["use_wham_noise"]:
        sb.dataio.dataset.set_output_keys(
            datasets, ["id", "mix_sig", "s1_sig", "s2_sig", "noise_sig"]
        )
    elif (hparams["num_spks"] == 3) and hparams["use_wham_noise"]:
        sb.dataio.dataset.set_output_keys(
            datasets,
            ["id", "mix_sig", "s1_sig", "s2_sig", "s3_sig", "noise_sig"],
        )
    elif (hparams["num_spks"] == 2) and not hparams["use_wham_noise"]:
        sb.dataio.dataset.set_output_keys(
            datasets, ["id", "mix_sig", "s1_sig", "s2_sig"]
        )
    else:
        sb.dataio.dataset.set_output_keys(
            datasets, ["id", "mix_sig", "s1_sig", "s2_sig", "s3_sig"]
        )

    return train_data, valid_data, test_data


if __name__ == "__main__":

    # Load hyperparameters file with command-line overrides
    hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    # Initialize ddp (useful only for multi-GPU DDP training)
    sb.utils.distributed.ddp_init_group(run_opts)

    # Logger info
    logger = logging.getLogger(__name__)

    # Create experiment directory
    sb.create_experiment_directory(
        experiment_directory=hparams["output_folder"],
        hyperparams_to_save=hparams_file,
        overrides=overrides,
    )

    # Check if wsj0_tr is set with dynamic mixing
    if hparams["dynamic_mixing"] and not os.path.exists(
            hparams["base_folder_dm"]
    ):
        print(
            "Please, specify a valid base_folder_dm folder when using dynamic mixing"
        )
        sys.exit(1)

    # Data preparation
    from prepare_data import prepare_librimix

    run_on_main(
        prepare_librimix,
        kwargs={
            "datapath": hparams["data_folder"],
            "savepath": hparams["save_folder"],
            "n_spks": hparams["num_spks"],
            "skip_prep": hparams["skip_prep"],
            "librimix_addnoise": hparams["use_wham_noise"],
            "fs": hparams["sample_rate"],
        },
    )

    # Create dataset objects
    if hparams["dynamic_mixing"]:
        from dynamic_mixing import (
            dynamic_mix_data_prep_librimix as dynamic_mix_data_prep,
        )

        # if the base_folder for dm is not processed, preprocess them
        if "processed" not in hparams["base_folder_dm"]:
            # if the processed folder already exists we just use it otherwise we do the preprocessing
            if not os.path.exists(
                    os.path.normpath(hparams["base_folder_dm"]) + "_processed"
            ):
                from preprocess_dynamic_mixing import (
                    resample_folder,
                )

                print("Resampling the base folder")
                run_on_main(
                    resample_folder,
                    kwargs={
                        "input_folder": hparams["base_folder_dm"],
                        "output_folder": os.path.normpath(
                            hparams["base_folder_dm"]
                        )
                                         + "_processed",
                        "fs": hparams["sample_rate"],
                        "regex": "**/*.flac",
                    },
                )
                # adjust the base_folder_dm path
                hparams["base_folder_dm"] = (
                        os.path.normpath(hparams["base_folder_dm"]) + "_processed"
                )
            else:
                print(
                    "Using the existing processed folder on the same directory as base_folder_dm"
                )
                hparams["base_folder_dm"] = (
                        os.path.normpath(hparams["base_folder_dm"]) + "_processed"
                )

        dm_hparams = {
            "train_data": hparams["train_data"],
            "data_folder": hparams["data_folder"],
            "base_folder_dm": hparams["base_folder_dm"],
            "sample_rate": hparams["sample_rate"],
            "num_spks": hparams["num_spks"],
            "training_signal_len": hparams["training_signal_len"],
            "dataloader_opts": hparams["dataloader_opts"],
            "use_wham_noise": hparams["use_wham_noise"],
        }

        train_data = dynamic_mix_data_prep(dm_hparams)
        _, valid_data, test_data = dataio_prep(hparams)
    else:
        train_data, valid_data, test_data = dataio_prep(hparams)

    # Load pretrained model if pretrained_separator is present in the yaml
    if "pretrained_separator" in hparams:
        run_on_main(hparams["pretrained_separator"].collect_files)
        hparams["pretrained_separator"].load_collected()

    # Brain class initialization
    separator = Separation(
        modules=hparams["modules"],
        opt_class=hparams["optimizer"],
        hparams=hparams,
        run_opts=run_opts,
        checkpointer=hparams["checkpointer"],
    )

    # re-initialize the parameters if we don't use a pretrained model
    if "pretrained_separator" not in hparams:
        for module in separator.modules.values():
            separator.reset_layer_recursively(module)

    if not hparams["test_only"]:
        # Training
        separator.fit(
            separator.hparams.epoch_counter,
            train_data,
            valid_data,
            train_loader_kwargs=hparams["dataloader_opts"],
            valid_loader_kwargs=hparams["dataloader_opts"],
        )

    # Eval
    separator.evaluate(test_data, min_key="si-snr")
    separator.save_results(test_data)

recipes/LibriMix/separation/train.py
TzuchengChang/NASS
"""Core SpeechBrain code for running experiments.

Authors
 * Peter Plantinga 2020
 * Abdel Heba 2020
 * Mirco Ravanelli 2020
 * Aku Rouhe 2021
 * Andreas Nautsch 2022
"""

import os
import sys
import yaml
import time
import torch
import shutil
import logging
import inspect
import pathlib
import argparse
import tempfile
import warnings
from contextlib import contextmanager
import speechbrain as sb
from datetime import date
from enum import Enum, auto
from tqdm.contrib import tqdm
from types import SimpleNamespace
from torch.nn import SyncBatchNorm
from torch.utils.data import DataLoader
from torch.nn import DataParallel as DP
from torch.utils.data import IterableDataset
from torch.utils.data import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from hyperpyyaml import resolve_references
from speechbrain.utils.distributed import run_on_main
from speechbrain.utils.optimizers import rm_vector_weight_decay
from speechbrain.dataio.dataloader import LoopedLoader
from speechbrain.dataio.dataloader import SaveableDataLoader
from speechbrain.dataio.sampler import DistributedSamplerWrapper
from speechbrain.dataio.sampler import ReproducibleRandomSampler

logger = logging.getLogger(__name__)
DEFAULT_LOG_CONFIG = os.path.dirname(os.path.abspath(__file__))
DEFAULT_LOG_CONFIG = os.path.join(DEFAULT_LOG_CONFIG, "log-config.yaml")
torch._C._jit_set_profiling_executor(False)
torch._C._jit_set_profiling_mode(False)
INTRA_EPOCH_CKPT_FLAG = "brain_intra_epoch_ckpt"
PYTHON_VERSION_MAJOR = 3
PYTHON_VERSION_MINOR = 7


def create_experiment_directory(
        experiment_directory,
        hyperparams_to_save=None,
        overrides={},
        log_config=DEFAULT_LOG_CONFIG,
        save_env_desc=True,
):
    """Create the output folder and relevant experimental files.

    Arguments
    ---------
    experiment_directory : str
        The place where the experiment directory should be created.
    hyperparams_to_save : str
        A filename of a yaml file representing the parameters for this
        experiment. If passed, references are resolved, and the result is
        written to a file in the experiment directory called "hyperparams.yaml".
    overrides : dict
        A mapping of replacements made in the yaml file, to save in yaml.
    log_config : str
        A yaml filename containing configuration options for the logger.
    save_env_desc : bool
        If True, an environment state description is saved to the experiment
        directory, in a file called env.log in the experiment directory.
    """
    try:
        # all writing command must be done with the main_process
        if sb.utils.distributed.if_main_process():
            if not os.path.isdir(experiment_directory):
                os.makedirs(experiment_directory)

            # Write the parameters file
            if hyperparams_to_save is not None:
                hyperparams_filename = os.path.join(
                    experiment_directory, "hyperparams.yaml"
                )
                with open(hyperparams_to_save) as f:
                    resolved_yaml = resolve_references(f, overrides)
                with open(hyperparams_filename, "w") as w:
                    print("# Generated %s from:" % date.today(), file=w)
                    print("# %s" % os.path.abspath(hyperparams_to_save), file=w)
                    print("# yamllint disable", file=w)
                    shutil.copyfileobj(resolved_yaml, w)

            # Copy executing file to output directory
            module = inspect.getmodule(inspect.currentframe().f_back)
            if module is not None:
                callingfile = os.path.realpath(module.__file__)
                shutil.copy(callingfile, experiment_directory)

            # Log exceptions to output automatically
            log_file = os.path.join(experiment_directory, "log.txt")
            logger_overrides = {
                "handlers": {"file_handler": {"filename": log_file}}
            }
            sb.utils.logger.setup_logging(log_config, logger_overrides)
            sys.excepthook = _logging_excepthook

            # Log beginning of experiment!
            logger.info("Beginning experiment!")
            logger.info(f"Experiment folder: {experiment_directory}")

            # Save system description:
            if save_env_desc:
                description_str = sb.utils.logger.get_environment_description()
                with open(
                        os.path.join(experiment_directory, "env.log"), "w"
                ) as fo:
                    fo.write(description_str)
    finally:
        # wait for main_process if ddp is used
        sb.utils.distributed.ddp_barrier()


def _logging_excepthook(exc_type, exc_value, exc_traceback):
    """Interrupt exception raising to log the error."""
    logger.error("Exception:", exc_info=(exc_type, exc_value, exc_traceback))


def parse_arguments(arg_list=None):
    """Parse command-line arguments to the experiment.

    Arguments
    ---------
    arg_list : list, None
        A list of arguments to parse.  If not given, this is read from
        `sys.argv[1:]`

    Returns
    -------
    param_file : str
        The location of the parameters file.
    run_opts : dict
        Run options, such as distributed, device, etc.
    overrides : dict
        The overrides to pass to ``load_hyperpyyaml``.

    Example
    -------
    >>> argv = ['hyperparams.yaml', '--device', 'cuda:1', '--seed', '10']
    >>> filename, run_opts, overrides = parse_arguments(argv)
    >>> filename
    'hyperparams.yaml'
    >>> run_opts["device"]
    'cuda:1'
    >>> overrides
    'seed: 10'
    """
    if arg_list is None:
        arg_list = sys.argv[1:]
    parser = argparse.ArgumentParser(description="Run a SpeechBrain experiment")
    parser.add_argument(
        "param_file",
        type=str,
        help="A yaml-formatted file using the extended YAML syntax. "
             "defined by SpeechBrain.",
    )
    parser.add_argument(
        "--debug",
        default=False,
        action="store_true",
        help="Run the experiment with only a few batches for all "
             "datasets, to ensure code runs without crashing.",
    )
    parser.add_argument(
        "--debug_batches",
        type=int,
        default=2,
        help="Number of batches to run in debug mode.",
    )
    parser.add_argument(
        "--debug_epochs",
        type=int,
        default=2,
        help="Number of epochs to run in debug mode. "
             "If a non-positive number is passed, all epochs are run.",
    )
    parser.add_argument(
        "--debug_persistently",
        default=False,
        action="store_true",
        help="Keep data stored during debug mode (not using /tmp).",
    )
    parser.add_argument(
        "--log_config",
        type=str,
        help="A file storing the configuration options for logging",
    )
    # if use_env = False in torch.distributed.lunch then local_rank arg is given
    parser.add_argument(
        "--local_rank",
        "--local-rank",  # alias required for PyTorch 2.x
        type=int,
        help="Rank on local machine",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="The device to run the experiment on (e.g. 'cuda:0')",
    )
    parser.add_argument(
        "--data_parallel_backend",
        default=False,
        action="store_true",
        help="This flag enables training with data_parallel.",
    )
    parser.add_argument(
        "--distributed_launch",
        default=False,
        action="store_true",
        help="This flag enables training with DDP. Assumes script run with "
             "`torch.distributed.launch`",
    )
    parser.add_argument(
        "--distributed_backend",
        type=str,
        default="nccl",
        help="One of {nccl, gloo, mpi}",
    )
    parser.add_argument(
        "--find_unused_parameters",
        default=False,
        action="store_true",
        help="This flag disable unused parameters detection",
    )
    parser.add_argument(
        "--jit_module_keys",
        type=str,
        nargs="*",
        help="A list of keys in the 'modules' dict to jitify",
    )
    parser.add_argument(
        "--auto_mix_prec",
        default=None,
        action="store_true",
        help="This flag enables training with automatic mixed-precision.",
    )
    parser.add_argument(
        "--bfloat16_mix_prec",
        default=None,
        action="store_true",
        help="This flag enables training with bfloat16 mixed-precision.",
    )
    parser.add_argument(
        "--max_grad_norm",
        type=float,
        help="Gradient norm will be clipped to this value, "
             "enter negative value to disable.",
    )
    parser.add_argument(
        "--nonfinite_patience",
        type=int,
        help="Max number of batches per epoch to skip if loss is nonfinite.",
    )
    parser.add_argument(
        "--noprogressbar",
        default=None,
        action="store_true",
        help="This flag disables the data loop progressbars.",
    )
    parser.add_argument(
        "--ckpt_interval_minutes",
        type=float,
        help="Amount of time between saving intra-epoch checkpoints "
             "in minutes. If non-positive, intra-epoch checkpoints are not saved.",
    )
    parser.add_argument(
        "--grad_accumulation_factor",
        type=int,
        help="Number of batches to accumulate gradients before optimizer step",
    )
    parser.add_argument(
        "--optimizer_step_limit",
        type=int,
        help="Number of optimizer steps to run. If not passed, all epochs are run.",
    )
    parser.add_argument(
        "--tqdm_colored_bar",
        default=False,
        action="store_true",
        help="Enable colored progress-bar in tqdm. If this is "
             "false, tqdm shall use default colors.",
    )
    parser.add_argument(
        "--remove_vector_weight_decay",
        default=False,
        action="store_true",
        help="Make vectors (e.g. norms and biases) a separate parameter group without weight_decay.",
    )

    # Accept extra args to override yaml
    run_opts, overrides = parser.parse_known_args(arg_list)

    # Ignore items that are "None", they were not passed
    run_opts = {k: v for k, v in vars(run_opts).items() if v is not None}

    param_file = run_opts["param_file"]
    del run_opts["param_file"]

    overrides = _convert_to_yaml(overrides)

    # Checking that DataParallel use the right number of GPU
    if run_opts["data_parallel_backend"]:
        if torch.cuda.device_count() == 0:
            raise ValueError("You must have at least 1 GPU.")

    # For DDP, the device args must equal to local_rank used by
    # torch.distributed.launch. If run_opts["local_rank"] exists,
    # use os.environ["LOCAL_RANK"]
    local_rank = None
    if "local_rank" in run_opts:
        local_rank = run_opts["local_rank"]
    else:
        if "LOCAL_RANK" in os.environ and os.environ["LOCAL_RANK"] != "":
            local_rank = int(os.environ["LOCAL_RANK"])

    # force device arg to be the same as local_rank from torch.distributed.lunch
    if local_rank is not None and "cuda" in run_opts["device"]:
        run_opts["device"] = run_opts["device"][:-1] + str(local_rank)

    return param_file, run_opts, overrides


def _convert_to_yaml(overrides):
    """Convert args to yaml for overrides"""
    yaml_string = ""

    # Handle '--arg=val' type args
    joined_args = "=".join(overrides)
    split_args = joined_args.split("=")

    for arg in split_args:
        if arg.startswith("--"):
            yaml_string += "\n" + arg[len("--"):] + ":"
        else:
            yaml_string += " " + arg

    return yaml_string.strip()


class Stage(Enum):
    """Simple enum to track stage of experiments."""

    TRAIN = auto()
    VALID = auto()
    TEST = auto()


@sb.utils.checkpoints.register_checkpoint_hooks
class Brain:
    """Brain class abstracts away the details of data loops.

    The primary purpose of the `Brain` class is the implementation of
    the ``fit()`` method, which iterates epochs and datasets for the
    purpose of "fitting" a set of modules to a set of data.

    In order to use the ``fit()`` method, one should sub-class the ``Brain``
    class and override any methods for which the default behavior does not
    match the use case. For a simple use case (e.g., training a single model
    with a single dataset) the only methods that need to be overridden are:

    * ``compute_forward()``
    * ``compute_objectives()``

    The example below illustrates how overriding these two methods is done.

    For more complicated use cases, such as multiple modules that need to
    be updated, the following methods can be overridden:

    * ``fit_batch()``
    * ``evaluate_batch()``

    Arguments
    ---------
    modules : dict of str:torch.nn.Module pairs
        These modules are passed to the optimizer by default if they have
        trainable parameters, and will have ``train()``/``eval()`` called on them.
    opt_class : torch.optim class
        A torch optimizer constructor that takes only the list of
        parameters (e.g. a lambda or partial function definition). By default,
        this will be passed all modules in ``modules`` at the
        beginning of the ``fit()`` method. This behavior can be changed
        by overriding the ``configure_optimizers()`` method.
    hparams : dict
        Each key:value pair should consist of a string key and a hyperparameter
        that is used within the overridden methods. These will
        be accessible via an ``hparams`` attribute, using "dot" notation:
        e.g., self.hparams.model(x).
    run_opts : dict
        A set of options to change the runtime environment, including

        debug (bool)
            If ``True``, this will only iterate a few batches for all
            datasets, to ensure code runs without crashing.
        debug_batches (int)
            Number of batches to run in debug mode, Default ``2``.
        debug_epochs (int)
            Number of epochs to run in debug mode, Default ``2``.
            If a non-positive number is passed, all epochs are run.
        debug_persistently (bool)
            Keep data stored during debug mode (not using /tmp), Default ``False``.
        jit_module_keys (list of str)
            List of keys in ``modules`` that should be jit compiled.
        distributed_backend (str)
            One of ``nccl``, ``gloo``, ``mpi``.
        device (str)
            The location for performing computations.
        auto_mix_prec (bool)
            If ``True``, automatic mixed-precision is used.
            Activate it only with cuda.
        max_grad_norm (float)
            Default implementation of ``fit_batch()`` uses
            ``clip_grad_norm_`` with this value. Default: ``5``.
        nonfinite_patience (int)
            Number of times to ignore non-finite losses before stopping.
            Default: ``3``.
        noprogressbar (bool)
            Whether to turn off progressbar when training. Default: ``False``.
        ckpt_interval_minutes (float)
            Amount of time between saving intra-epoch checkpoints,
            in minutes, default: ``15.0``. If non-positive, these are not saved.

        Typically in a script this comes from ``speechbrain.parse_args``, which
        has different defaults than Brain. If an option is not defined here
        (keep in mind that parse_args will inject some options by default),
        then the option is also searched for in hparams (by key).
    checkpointer : speechbrain.Checkpointer
        By default, this will be used to load checkpoints, and will have the
        optimizer added to continue training if interrupted.
    profiler : torch.profiler.profile
        Context manager for profiling and benchmarking of training/inference steps.
        Default: ``None`` (skip profiling).

    Example
    -------
    >>> from torch.optim import SGD
    >>> class SimpleBrain(Brain):
    ...     def compute_forward(self, batch, stage):
    ...         return self.modules.model(batch[0])
    ...     def compute_objectives(self, predictions, batch, stage):
    ...         return torch.nn.functional.l1_loss(predictions, batch[0])
    >>> model = torch.nn.Linear(in_features=10, out_features=10)
    >>> brain = SimpleBrain({"model": model}, opt_class=lambda x: SGD(x, 0.1))
    >>> brain.fit(range(1), ([torch.rand(10, 10), torch.rand(10, 10)],))
    """

    def __init__(  # noqa: C901
            self,
            modules=None,
            opt_class=None,
            hparams=None,
            run_opts=None,
            checkpointer=None,
            profiler=None,
    ):
        self.opt_class = opt_class
        self.checkpointer = checkpointer
        self.profiler = profiler

        # Arguments passed via the run opts dictionary
        run_opt_defaults = {
            "debug": False,
            "debug_batches": 2,
            "debug_epochs": 2,
            "debug_persistently": False,
            "device": "cpu",
            "data_parallel_backend": False,
            "distributed_launch": False,
            "distributed_backend": "nccl",
            "find_unused_parameters": False,
            "jit_module_keys": None,
            "auto_mix_prec": False,
            "bfloat16_mix_prec": False,
            "max_grad_norm": 5.0,
            "nonfinite_patience": 3,
            "noprogressbar": False,
            "ckpt_interval_minutes": 0,
            "grad_accumulation_factor": 1,
            "optimizer_step_limit": None,
            "tqdm_colored_bar": False,
            "tqdm_barcolor": {
                "train": "GREEN",
                "valid": "MAGENTA",
                "test": "CYAN",
            },
            "remove_vector_weight_decay": False,
        }

        for arg, default in run_opt_defaults.items():
            if run_opts is not None and arg in run_opts:
                if hparams is not None and arg in hparams:
                    logger.info(
                        "Info: "
                        + arg
                        + " arg overridden by command line input to: "
                        + str(run_opts[arg])
                    )
                setattr(self, arg, run_opts[arg])
            else:
                # If any arg from run_opt_defaults exist in hparams and
                # not in command line args "run_opts"
                if hparams is not None and arg in hparams:
                    logger.info(
                        "Info: " + arg + " arg from hparam file is used"
                    )
                    setattr(self, arg, hparams[arg])
                else:
                    setattr(self, arg, default)

        # Check Python version
        if not (
                sys.version_info.major == PYTHON_VERSION_MAJOR
                and sys.version_info.minor >= PYTHON_VERSION_MINOR
        ):
            logger.warning(
                "Detected Python "
                + str(sys.version_info.major)
                + "."
                + str(sys.version_info.minor)
                + ". We suggest using SpeechBrain with Python >="
                + str(PYTHON_VERSION_MAJOR)
                + "."
                + str(PYTHON_VERSION_MINOR)
            )

        if self.data_parallel_backend and self.distributed_launch:
            sys.exit(
                "To use data_parallel backend, start your script with:\n\t"
                "python experiment.py hyperparams.yaml "
                "--data_parallel_backend=True"
                "To use DDP backend, start your script with:\n\t"
                "python -m torch.distributed.lunch [args]\n"
                "experiment.py hyperparams.yaml --distributed_launch=True "
                "--distributed_backend=nccl"
            )

        # Switch to the right context
        if self.device == "cuda":
            torch.cuda.set_device(0)
        elif "cuda" in self.device:
            torch.cuda.set_device(int(self.device[-1]))

        # Put modules on the right device, accessible with dot notation
        self.modules = torch.nn.ModuleDict(modules).to(self.device)

        # Make hyperparams available with dot notation too
        if hparams is not None:
            self.hparams = SimpleNamespace(**hparams)

        # Checkpointer should point at a temporary directory in debug mode
        if (
                self.debug
                and not self.debug_persistently
                and self.checkpointer is not None
                and hasattr(self.checkpointer, "checkpoints_dir")
        ):
            tempdir = tempfile.TemporaryDirectory()
            logger.info(
                "Since debug mode is active, switching checkpointer "
                f"output to temporary directory: {tempdir.name}"
            )
            self.checkpointer.checkpoints_dir = pathlib.Path(tempdir.name)

            # Keep reference to tempdir as long as checkpointer exists
            self.checkpointer.tempdir = tempdir

        # Sampler should be handled by `make_dataloader`
        # or if you provide a DataLoader directly, you can set
        # this.train_sampler = your_sampler
        # to have your_sampler.set_epoch() called on each epoch.
        self.train_sampler = None

        # Automatic mixed precision init
        if self.auto_mix_prec:
            self.scaler = torch.cuda.amp.GradScaler()
            if self.checkpointer is not None:
                self.checkpointer.add_recoverable("scaler", self.scaler)

        # List parameter count for the user
        total_params = sum(
            p.numel() for p in self.modules.parameters() if p.requires_grad
        )
        if total_params > 0:
            clsname = self.__class__.__name__
            fmt_num = sb.utils.logger.format_order_of_magnitude(total_params)
            logger.info(f"{fmt_num} trainable parameters in {clsname}")

        if self.distributed_launch:
            self.rank = int(os.environ["RANK"])
            if not torch.distributed.is_initialized():
                if self.rank > 0:
                    sys.exit(
                        " ================ WARNING ==============="
                        "Please add sb.ddp_init_group() into your exp.py"
                        "To use DDP backend, start your script with:\n\t"
                        "python -m torch.distributed.launch [args]\n\t"
                        "experiment.py hyperparams.yaml "
                        "--distributed_launch=True --distributed_backend=nccl"
                    )
                else:
                    logger.warning(
                        "To use DDP, please add "
                        "sb.utils.distributed.ddp_init_group() into your exp.py"
                    )
                    logger.info(
                        "Only the main process is alive, "
                        "all other subprocess were killed."
                    )

        # Prepare iterating variables
        self.avg_train_loss = 0.0
        self.step = 0
        self.optimizer_step = 0

        # Add this class to the checkpointer for intra-epoch checkpoints
        if self.checkpointer is not None:
            self.checkpointer.add_recoverable("brain", self)

        # Force default color for tqdm progrressbar
        if not self.tqdm_colored_bar:
            self.tqdm_barcolor = dict.fromkeys(self.tqdm_barcolor, "")

    def compute_forward(self, batch, stage):
        """Forward pass, to be overridden by sub-classes.

        Arguments
        ---------
        batch : torch.Tensor or tensors
            An element from the dataloader, including inputs for processing.
        stage : Stage
            The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

        Returns
        -------
        torch.Tensor or Tensors
            The outputs after all processing is complete.
            Directly passed to ``compute_objectives()``.
        """
        raise NotImplementedError

    def compute_objectives(self, predictions, batch, stage):
        """Compute loss, to be overridden by sub-classes.

        Arguments
        ---------
        predictions : torch.Tensor or Tensors
            The output tensor or tensors to evaluate.
            Comes directly from ``compute_forward()``.
        batch : torch.Tensor or tensors
            An element from the dataloader, including targets for comparison.
        stage : Stage
            The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST

        Returns
        -------
        loss : torch.Tensor
            A tensor with the computed loss.
        """
        raise NotImplementedError

    def on_stage_start(self, stage, epoch=None):
        """Gets called when a stage starts.

        Useful for defining class variables used during the stage.

        Arguments
        ---------
        stage : Stage
            The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST
        epoch : int
            The current epoch count.
        """
        pass

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage.

        Useful for computing stage statistics, saving checkpoints, etc.

        Arguments
        ---------
        stage : Stage
            The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST
        stage_loss : float
            The average loss over the completed stage.
        epoch : int
            The current epoch count.
        """
        pass

    def make_dataloader(
            self, dataset, stage, ckpt_prefix="dataloader-", **loader_kwargs
    ):
        """Creates DataLoaders for Datasets.

        This is used by ``fit()`` and ``evaluate()`` if they just receive
        Datasets.

        Alternatively, this can be called from outside the Brain subclass.
        In that case, the DataLoader should be passed to ``fit()`` in place
        of the dataset.

        The Stage.TRAIN DataLoader is handled specially. It has extra args for
        shuffle and drop_last. In DDP a DistributedSampler is created (unless
        the dataset is an IterableDataset).

        NOTE
        ----
        Some important DataLoader arguments are passed via **loader_kwargs,
        e.g., batch_size, num_workers, pin_memory.

        NOTE
        ----
        By default, ``evaluate()`` specifies ckpt_prefix=None to stop the test
        DataLoader being added to the checkpointer. If you need to add a
        recoverable after saving checkpoints (e.g., at test time, after
        checkpointing the training), and still be able to recover reasonably,
        you should probably specify ``allow_partial_load=True``.

        Arguments
        ---------
        dataset : Dataset
            A set of data to use to create data loader. If the Dataset is a
            DynamicItemDataset, PaddedBatch is used as the default collate_fn,
            unless specified in loader_kwargs.
        stage : Stage
            The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST
        ckpt_prefix : str, None
            Prefix to use for SaveableDataLoader Checkpoint name. The Stage
            name is added to this to create the full key. Set to None to not
            save the DataLoader.
        **loader_kwargs : dict
            Additional keyword arguments to the DataLoader.
            E.g., batch_size, num_workers, pin_memory.
        """
        # TRAIN stage is handled specially.
        if stage == sb.Stage.TRAIN:
            loader_kwargs = self._train_loader_specifics(dataset, loader_kwargs)
        # This commented-out code block is useful when one can ensure
        # metric reporting is DDP-valid for VALID & EVAL datasets.
        # elif self.distributed_launch:
        #     loader_kwargs = sb.dataio.dataloader.distributed_loader_specifics(
        #         self.distributed_launch, self.rank, dataset, loader_kwargs
        #     )
        dataloader = sb.dataio.dataloader.make_dataloader(
            dataset, **loader_kwargs
        )

        if (
                self.checkpointer is not None
                and ckpt_prefix is not None
                and (
                isinstance(dataloader, SaveableDataLoader)
                or isinstance(dataloader, LoopedLoader)
        )
        ):
            ckpt_key = ckpt_prefix + stage.name
            self.checkpointer.add_recoverable(ckpt_key, dataloader)
        return dataloader

    def _train_loader_specifics(self, dataset, loader_kwargs):
        sampler = loader_kwargs.get("sampler", None)
        # Shuffling should really only matter for the train stage. Shuffling
        # will also lead to more padding in batches if the order was otherwise
        # sorted by length.
        shuffle = loader_kwargs.get("shuffle", False)
        if shuffle and not self.distributed_launch:
            if sampler is not None:
                raise ValueError(
                    "Cannot specify both shuffle=True"
                    "and a sampler in loader_kwargs"
                )
            sampler = ReproducibleRandomSampler(dataset)
            self.train_sampler = sampler
            loader_kwargs["sampler"] = self.train_sampler
            # Delete the shuffle flag, since you cannot specify both a sampler and
            # shuffling:
            del loader_kwargs["shuffle"]

        # Possibly make a DistributedSampler or a wrapper for some other sampler
        if self.distributed_launch and not isinstance(dataset, IterableDataset):
            # sort or not
            if hasattr(self.hparams, "sorting"):
                shuffle_ddp = (
                        self.hparams.sorting == "random"
                )  # False if 'ascending' or 'descending'
            else:
                shuffle_ddp = True

            drop_last = loader_kwargs.get("drop_last", False)
            # num_replicas arg is equal to world_size
            # and retrieved automatically within
            # DistributedSampler obj.
            if sampler is not None:
                self.train_sampler = DistributedSamplerWrapper(
                    sampler,
                    rank=self.rank,
                    drop_last=drop_last,
                    shuffle=shuffle,
                )

                # with DistributedSamplerWrapper, one must disable shuffling for dataloader
                loader_kwargs["shuffle"] = False
                loader_kwargs["sampler"] = self.train_sampler
            elif loader_kwargs.get("batch_sampler") is None:
                # no sampler and batch-sampler
                self.train_sampler = DistributedSampler(
                    dataset,
                    rank=self.rank,
                    shuffle=shuffle_ddp,
                    drop_last=drop_last,
                )

                # with DistributedSamplerWrapper, one must disable shuffling for dataloader
                loader_kwargs["shuffle"] = False
                loader_kwargs["sampler"] = self.train_sampler
            else:  # batch_sampler was specified
                self.train_sampler = DistributedSamplerWrapper(
                    loader_kwargs.get("batch_sampler", None),
                    rank=self.rank,
                    shuffle=shuffle_ddp,
                )
                loader_kwargs["batch_sampler"] = self.train_sampler
        elif self.distributed_launch and isinstance(dataset, IterableDataset):
            logger.warning(
                "Cannot automatically solve distributed sampling "
                "for IterableDataset."
            )
        return loader_kwargs

    def on_fit_start(self):
        """Gets called at the beginning of ``fit()``, on multiple processes
        if ``distributed_count > 0`` and backend is ddp.

        Default implementation compiles the jit modules, initializes
        optimizers, and loads the latest checkpoint to resume training.
        """
        # Run this *after* starting all processes since jit modules cannot be
        # pickled.
        self._compile_jit()

        # Wrap modules with parallel backend after jit
        self._wrap_distributed()

        # Initialize optimizers after parameters are configured
        self.init_optimizers()

        # Load latest checkpoint to resume training if interrupted
        if self.checkpointer is not None:
            self.checkpointer.recover_if_possible(
                device=torch.device(self.device)
            )

    def init_optimizers(self):
        """Called during ``on_fit_start()``, initialize optimizers
        after parameters are fully configured (e.g. DDP, jit).

        The default implementation of this method depends on an optimizer
        class being passed at initialization that takes only a list
        of parameters (e.g., a lambda or a partial function definition).
        This creates a single optimizer that optimizes all trainable params.

        Override this class if there are multiple optimizers.
        """

        all_params = self.modules.parameters()

        if self.opt_class is not None:
            if self.remove_vector_weight_decay:
                all_params = rm_vector_weight_decay(self.modules)

            self.optimizer = self.opt_class(all_params)

            if self.checkpointer is not None:
                self.checkpointer.add_recoverable("optimizer", self.optimizer)

    def zero_grad(self, set_to_none=False):
        """Sets the gradients of all optimized ``torch.Tensor``s to zero
        if ``set_to_none=False`` (default) or to None otherwise.

        Setting gradients to None should save the memory, e.g.
        during ``evaluate()`` and thus larger batch might be used.
        """
        if hasattr(self, "optimizer"):
            self.optimizer.zero_grad(set_to_none)

    def on_evaluate_start(self, max_key=None, min_key=None):
        """Gets called at the beginning of ``evaluate()``

        Default implementation loads the best-performing checkpoint for
        evaluation, based on stored metrics.

        Arguments
        ---------
        max_key : str
            Key to use for finding best checkpoint (higher is better).
            By default, passed to ``self.checkpointer.recover_if_possible()``.
        min_key : str
            Key to use for finding best checkpoint (lower is better).
            By default, passed to ``self.checkpointer.recover_if_possible()``.
        """

        # Recover best checkpoint for evaluation
        if self.checkpointer is not None:
            self.checkpointer.recover_if_possible(
                max_key=max_key,
                min_key=min_key,
                device=torch.device(self.device),
            )

    def fit_batch(self, batch):
        """Fit one batch, override to do multiple updates.

        The default implementation depends on a few methods being defined
        with a particular behavior:

        * ``compute_forward()``
        * ``compute_objectives()``

        Also depends on having optimizers passed at initialization.

        Arguments
        ---------
        batch : list of torch.Tensors
            Batch of data to use for training. Default implementation assumes
            this batch has two elements: inputs and targets.

        Returns
        -------
        detached loss
        """
        should_step = self.step % self.grad_accumulation_factor == 0
        # Managing automatic mixed precision
        if self.auto_mix_prec:
            with torch.autocast(device_type=torch.device(self.device).type):
                outputs = self.compute_forward(batch, Stage.TRAIN)

            # Losses are excluded from mixed precision to avoid instabilities
            loss = self.compute_objectives(outputs, batch, Stage.TRAIN)
            with self.no_sync(not should_step):
                self.scaler.scale(
                    loss / self.grad_accumulation_factor
                ).backward()
            if should_step:
                self.scaler.unscale_(self.optimizer)
                if self.check_gradients(loss):
                    self.scaler.step(self.optimizer)
                self.scaler.update()
                self.zero_grad()
                self.optimizer_step += 1
        else:
            if self.bfloat16_mix_prec:
                with torch.autocast(
                        device_type=torch.device(self.device).type,
                        dtype=torch.bfloat16,
                ):
                    outputs = self.compute_forward(batch, Stage.TRAIN)
                    loss = self.compute_objectives(outputs, batch, Stage.TRAIN)
            else:
                outputs = self.compute_forward(batch, Stage.TRAIN)
                loss = self.compute_objectives(outputs, batch, Stage.TRAIN)

            with self.no_sync(not should_step):
                (loss / self.grad_accumulation_factor).backward()
            if should_step:
                if self.check_gradients(loss):
                    self.optimizer.step()
                self.zero_grad()
                self.optimizer_step += 1

        self.on_fit_batch_end(batch, outputs, loss, should_step)
        return loss.detach().cpu()

    def on_fit_batch_end(self, batch, outputs, loss, should_step):
        """Called after ``fit_batch()``, meant for calculating and logging metrics.

        Arguments
        ---------
        batch : list of torch.Tensors
            Batch of data to use for training. Default implementation assumes
            this batch has two elements: inputs and targets.
        outputs : list or dictionary of torch.Tensors
            Returned value of compute_forward().
        loss : torch.Tensor
            Returned value of compute_objectives().
        should_step : boolean
            Whether optimizer.step() was called or not.
        """
        pass

    def check_gradients(self, loss):
        """Check if gradients are finite and not too large.

        Automatically clips large gradients.

        Arguments
        ---------
        loss : tensor
            The loss tensor after ``backward()`` has been called but
            before the optimizers ``step()``.

        Returns
        -------
        bool
            Whether or not the optimizer step should be carried out.
        """
        if not torch.isfinite(loss):
            self.nonfinite_count += 1

            # Print helpful debug info
            logger.warning(f"Loss is {loss}.")
            for p in self.modules.parameters():
                if not torch.isfinite(p).all():
                    logger.warning("Parameter is not finite: " + str(p))

            # Check if patience is exhausted
            if self.nonfinite_count > self.nonfinite_patience:
                raise ValueError(
                    "Loss is not finite and patience is exhausted. "
                    "To debug, wrap `fit()` with "
                    "autograd's `detect_anomaly()`, e.g.\n\nwith "
                    "torch.autograd.detect_anomaly():\n\tbrain.fit(...)"
                )
            else:
                logger.warning(
                    "Patience not yet exhausted, ignoring this batch."
                )
                return False

        # Clip gradient norm
        if self.max_grad_norm > 0.0:
            torch.nn.utils.clip_grad_norm_(
                (p for p in self.modules.parameters()), self.max_grad_norm
            )

        return True

    def evaluate_batch(self, batch, stage):
        """Evaluate one batch, override for different procedure than train.

        The default implementation depends on two methods being defined
        with a particular behavior:

        * ``compute_forward()``
        * ``compute_objectives()``

        Arguments
        ---------
        batch : list of torch.Tensors
            Batch of data to use for evaluation. Default implementation assumes
            this batch has two elements: inputs and targets.
        stage : Stage
            The stage of the experiment: Stage.VALID, Stage.TEST

        Returns
        -------
        detached loss
        """

        out = self.compute_forward(batch, stage=stage)
        loss = self.compute_objectives(out, batch, stage=stage)
        return loss.detach().cpu()

    def _fit_train(self, train_set, epoch, enable):
        # Training stage
        self.on_stage_start(Stage.TRAIN, epoch)
        self.modules.train()
        self.zero_grad()

        # Reset nonfinite count to 0 each epoch
        self.nonfinite_count = 0

        if self.train_sampler is not None and hasattr(
                self.train_sampler, "set_epoch"
        ):
            self.train_sampler.set_epoch(epoch)

        # Time since last intra-epoch checkpoint
        last_ckpt_time = time.time()
        with tqdm(
                train_set,
                initial=self.step,
                dynamic_ncols=True,
                disable=not enable,
                colour=self.tqdm_barcolor["train"],
        ) as t:
            for batch in t:
                if self._optimizer_step_limit_exceeded:
                    logger.info("Train iteration limit exceeded")
                    break
                self.step += 1
                loss = self.fit_batch(batch)
                self.avg_train_loss = self.update_average(
                    loss, self.avg_train_loss
                )
                t.set_postfix(train_loss=self.avg_train_loss)

                # Profile only if desired (steps allow the profiler to know when all is warmed up)
                if self.profiler is not None:
                    if self.profiler.record_steps:
                        self.profiler.step()

                # Debug mode only runs a few batches
                if self.debug and self.step == self.debug_batches:
                    break

                if (
                        self.checkpointer is not None
                        and self.ckpt_interval_minutes > 0
                        and time.time() - last_ckpt_time
                        >= self.ckpt_interval_minutes * 60.0
                ):
                    # This should not use run_on_main, because that
                    # includes a DDP barrier. That eventually leads to a
                    # crash when the processes'
                    # time.time() - last_ckpt_time differ and some
                    # processes enter this block while others don't,
                    # missing the barrier.
                    if sb.utils.distributed.if_main_process():
                        self._save_intra_epoch_ckpt()
                    last_ckpt_time = time.time()

        # Run train "on_stage_end" on all processes
        self.zero_grad(set_to_none=True)  # flush gradients
        self.on_stage_end(Stage.TRAIN, self.avg_train_loss, epoch)
        self.avg_train_loss = 0.0
        self.step = 0

    def _fit_valid(self, valid_set, epoch, enable):
        # Validation stage
        if valid_set is not None:
            self.on_stage_start(Stage.VALID, epoch)
            self.modules.eval()
            avg_valid_loss = 0.0
            with torch.no_grad():
                for batch in tqdm(
                        valid_set,
                        dynamic_ncols=True,
                        disable=not enable,
                        colour=self.tqdm_barcolor["valid"],
                ):
                    self.step += 1
                    loss = self.evaluate_batch(batch, stage=Stage.VALID)
                    avg_valid_loss = self.update_average(loss, avg_valid_loss)

                    # Profile only if desired (steps allow the profiler to know when all is warmed up)
                    if self.profiler is not None:
                        if self.profiler.record_steps:
                            self.profiler.step()

                    # Debug mode only runs a few batches
                    if self.debug and self.step == self.debug_batches:
                        break

                # Only run validation "on_stage_end" on main process
                self.step = 0
                run_on_main(
                    self.on_stage_end,
                    args=[Stage.VALID, avg_valid_loss, epoch],
                )

    def fit(
            self,
            epoch_counter,
            train_set,
            valid_set=None,
            progressbar=None,
            train_loader_kwargs={},
            valid_loader_kwargs={},
    ):
        """Iterate epochs and datasets to improve objective.

        Relies on the existence of multiple functions that can (or should) be
        overridden. The following methods are used and expected to have a
        certain behavior:

        * ``fit_batch()``
        * ``evaluate_batch()``
        * ``update_average()``

        If the initialization was done with distributed_count > 0 and the
        distributed_backend is ddp, this will generally handle multiprocess
        logic, like splitting the training data into subsets for each device and
        only saving a checkpoint on the main process.

        Arguments
        ---------
        epoch_counter : iterable
            Each call should return an integer indicating the epoch count.
        train_set : Dataset, DataLoader
            A set of data to use for training. If a Dataset is given, a
            DataLoader is automatically created. If a DataLoader is given, it is
            used directly.
        valid_set : Dataset, DataLoader
            A set of data to use for validation. If a Dataset is given, a
            DataLoader is automatically created. If a DataLoader is given, it is
            used directly.
        train_loader_kwargs : dict
            Kwargs passed to `make_dataloader()` for making the train_loader
            (if train_set is a Dataset, not DataLoader).
            E.G. batch_size, num_workers.
            DataLoader kwargs are all valid.
        valid_loader_kwargs : dict
            Kwargs passed to `make_dataloader()` for making the valid_loader
            (if valid_set is a Dataset, not DataLoader).
            E.g., batch_size, num_workers.
            DataLoader kwargs are all valid.
        progressbar : bool
            Whether to display the progress of each epoch in a progressbar.
        """
        if not (
                isinstance(train_set, DataLoader)
                or isinstance(train_set, LoopedLoader)
        ):
            train_set = self.make_dataloader(
                train_set, stage=sb.Stage.TRAIN, **train_loader_kwargs
            )
        if valid_set is not None and not (
                isinstance(valid_set, DataLoader)
                or isinstance(valid_set, LoopedLoader)
        ):
            valid_set = self.make_dataloader(
                valid_set,
                stage=sb.Stage.VALID,
                ckpt_prefix=None,
                **valid_loader_kwargs,
            )

        self.on_fit_start()

        if progressbar is None:
            progressbar = not self.noprogressbar

        # Only show progressbar if requested and main_process
        enable = progressbar and sb.utils.distributed.if_main_process()

        # Iterate epochs
        for epoch in epoch_counter:
            self._fit_train(train_set=train_set, epoch=epoch, enable=enable)
            self._fit_valid(valid_set=valid_set, epoch=epoch, enable=enable)

            # Debug mode only runs a few epochs
            if (
                    self.debug
                    and epoch == self.debug_epochs
                    or self._optimizer_step_limit_exceeded
            ):
                break

    @property
    def _optimizer_step_limit_exceeded(self):
        return (
                self.optimizer_step_limit is not None
                and self.optimizer_step >= self.optimizer_step_limit
        )

    def _save_intra_epoch_ckpt(self):
        """Saves a CKPT with specific intra-epoch flag."""
        self.checkpointer.save_and_keep_only(
            end_of_epoch=False,
            num_to_keep=1,
            ckpt_predicate=lambda c: INTRA_EPOCH_CKPT_FLAG in c.meta,
            meta={INTRA_EPOCH_CKPT_FLAG: True},
            verbosity=logging.DEBUG,
        )

    def _compile_jit(self):
        """Compile requested modules with ``torch.jit.script``."""
        if self.jit_module_keys is None:
            return

        for name in self.jit_module_keys:
            if name not in self.modules:
                raise ValueError(
                    "module" + name + " is not defined in your hparams file."
                )
            module = torch.jit.script(self.modules[name])
            self.modules[name] = module.to(self.device)

    def _wrap_distributed(self):
        """Wrap modules with distributed wrapper when requested."""
        if not self.distributed_launch and not self.data_parallel_backend:
            return
        elif self.distributed_launch:
            for name, module in self.modules.items():
                if any(p.requires_grad for p in module.parameters()):
                    module = SyncBatchNorm.convert_sync_batchnorm(module)
                    if self.distributed_backend == "gloo":
                        module = DDP(
                            module,
                            device_ids=None,
                            find_unused_parameters=self.find_unused_parameters,
                        )
                    else:
                        module = DDP(
                            module,
                            device_ids=[self.device],
                            find_unused_parameters=self.find_unused_parameters,
                        )
                    self.modules[name] = module
        else:
            # data_parallel_backend
            for name, module in self.modules.items():
                if any(p.requires_grad for p in module.parameters()):
                    module = DP(module)
                    self.modules[name] = module

    def evaluate(
            self,
            test_set,
            max_key=None,
            min_key=None,
            progressbar=None,
            test_loader_kwargs={},
    ):
        """Iterate test_set and evaluate brain performance. By default, loads
        the best-performing checkpoint (as recorded using the checkpointer).

        Arguments
        ---------
        test_set : Dataset, DataLoader
            If a DataLoader is given, it is iterated directly. Otherwise passed
            to ``self.make_dataloader()``.
        max_key : str
            Key to use for finding best checkpoint, passed to
            ``on_evaluate_start()``.
        min_key : str
            Key to use for finding best checkpoint, passed to
            ``on_evaluate_start()``.
        progressbar : bool
            Whether to display the progress in a progressbar.
        test_loader_kwargs : dict
            Kwargs passed to ``make_dataloader()`` if ``test_set`` is not a
            DataLoader. NOTE: ``loader_kwargs["ckpt_prefix"]`` gets
            automatically overwritten to ``None`` (so that the test DataLoader
            is not added to the checkpointer).

        Returns
        -------
        average test loss
        """
        if progressbar is None:
            progressbar = not self.noprogressbar

        if not (
                isinstance(test_set, DataLoader)
                or isinstance(test_set, LoopedLoader)
        ):
            test_loader_kwargs["ckpt_prefix"] = None
            test_set = self.make_dataloader(
                test_set, Stage.TEST, **test_loader_kwargs
            )
        self.on_evaluate_start(max_key=max_key, min_key=min_key)
        self.on_stage_start(Stage.TEST, epoch=None)
        self.modules.eval()
        avg_test_loss = 0.0
        with torch.no_grad():
            for batch in tqdm(
                    test_set,
                    dynamic_ncols=True,
                    disable=not progressbar,
                    colour=self.tqdm_barcolor["test"],
            ):
                self.step += 1
                loss = self.evaluate_batch(batch, stage=Stage.TEST)
                avg_test_loss = self.update_average(loss, avg_test_loss)

                # Profile only if desired (steps allow the profiler to know when all is warmed up)
                if self.profiler is not None:
                    if self.profiler.record_steps:
                        self.profiler.step()

                # Debug mode only runs a few batches
                if self.debug and self.step == self.debug_batches:
                    break

            # Only run evaluation "on_stage_end" on main process
            run_on_main(
                self.on_stage_end, args=[Stage.TEST, avg_test_loss, None]
            )
        self.step = 0
        return avg_test_loss

    def update_average(self, loss, avg_loss):
        """Update running average of the loss.

        Arguments
        ---------
        loss : torch.tensor
            detached loss, a single float value.
        avg_loss : float
            current running average.

        Returns
        -------
        avg_loss : float
            The average loss.
        """
        if torch.isfinite(loss):
            avg_loss -= avg_loss / self.step
            avg_loss += float(loss) / self.step
        return avg_loss

    @contextmanager
    def no_sync(self, use=True):
        """Copies pytorch's implementation for doing no_sync across all modules.

        Explanation: nn.module.no_sync() is a context manager for when one does
        not want to sync gradients, which happens when using both DDP and gradient accumulation.
        Speechbrain brain's class can contain multiple modules and calling no_sync on these
        individually would be very awkward, therefore this contextmanager exists.

        Arguments
        ---------
        use : bool
            If set to `False` will still sync gradients, useful to make behaviour togglable.
        """
        if use:
            old_values_list = []
            for module in self.modules.values():
                if not hasattr(module, "require_backward_grad_sync"):
                    # if not using DDP
                    break
                old_values_list.append(module.require_backward_grad_sync)
                module.require_backward_grad_sync = False
            yield
            for module, old_value in zip(
                    self.modules.values(), old_values_list
            ):
                if not hasattr(module, "require_backward_grad_sync"):
                    break
                module.require_backward_grad_sync = old_value
        else:
            yield

    @sb.utils.checkpoints.mark_as_saver
    def _save(self, path):
        save_dict = {
            "step": self.step,
            "avg_train_loss": self.avg_train_loss,
            "optimizer_step": self.optimizer_step,
        }
        with open(path, "w") as w:
            w.write(yaml.dump(save_dict))

    @sb.utils.checkpoints.mark_as_loader
    def _recover(self, path, end_of_epoch, device):
        del end_of_epoch
        del device
        with open(path) as f:
            save_dict = yaml.safe_load(f)
        self.step = save_dict["step"]
        self.avg_train_loss = save_dict["avg_train_loss"]
        # Ensure compatibility with checkpoints from before optimizer_step:
        if "optimizer_step" not in save_dict:
            clsname = self.__class__.__name__
            MSG = f"'optimizer_step' not found in {clsname} checkpoint."
            MSG += " Using the saved 'step' value (BACKWARDS COMPATIBILITY)"
            warnings.warn(MSG)
            self.optimizer_step = self.step
        else:
            self.optimizer_step = save_dict["optimizer_step"]

speechbrain/core.py
TzuchengChang/NASS
"""Library to support dual-path speech separation.

Authors
 * Cem Subakan 2020
 * Mirco Ravanelli 2020
 * Samuele Cornell 2020
 * Mirko Bronzi 2020
 * Jianyuan Zhong 2020
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
from speechbrain.nnet.linear import Linear
from speechbrain.lobes.models.transformer.Transformer import TransformerEncoder
from speechbrain.lobes.models.transformer.Transformer import PositionalEncoding
from speechbrain.lobes.models.transformer.Conformer import ConformerEncoder
import speechbrain.nnet.RNN as SBRNN

from speechbrain.nnet.activations import Swish

EPS = 1e-8


class GlobalLayerNorm(nn.Module):
    """Calculate Global Layer Normalization.

    Arguments
    ---------
       dim : (int or list or torch.Size)
           Input shape from an expected input of size.
       eps : float
           A value added to the denominator for numerical stability.
       elementwise_affine : bool
          A boolean value that when set to True,
          this module has learnable per-element affine parameters
          initialized to ones (for weights) and zeros (for biases).

    Example
    -------
    >>> x = torch.randn(5, 10, 20)
    >>> GLN = GlobalLayerNorm(10, 3)
    >>> x_norm = GLN(x)
    """

    def __init__(self, dim, shape, eps=1e-8, elementwise_affine=True):
        super(GlobalLayerNorm, self).__init__()
        self.dim = dim
        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if self.elementwise_affine:
            if shape == 3:
                self.weight = nn.Parameter(torch.ones(self.dim, 1))
                self.bias = nn.Parameter(torch.zeros(self.dim, 1))
            if shape == 4:
                self.weight = nn.Parameter(torch.ones(self.dim, 1, 1))
                self.bias = nn.Parameter(torch.zeros(self.dim, 1, 1))
        else:
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)

    def forward(self, x):
        """Returns the normalized tensor.

        Arguments
        ---------
        x : torch.Tensor
            Tensor of size [N, C, K, S] or [N, C, L].
        """
        # x = N x C x K x S or N x C x L
        # N x 1 x 1
        # cln: mean,var N x 1 x K x S
        # gln: mean,var N x 1 x 1
        if x.dim() == 3:
            mean = torch.mean(x, (1, 2), keepdim=True)
            var = torch.mean((x - mean) ** 2, (1, 2), keepdim=True)
            if self.elementwise_affine:
                x = (
                        self.weight * (x - mean) / torch.sqrt(var + self.eps)
                        + self.bias
                )
            else:
                x = (x - mean) / torch.sqrt(var + self.eps)

        if x.dim() == 4:
            mean = torch.mean(x, (1, 2, 3), keepdim=True)
            var = torch.mean((x - mean) ** 2, (1, 2, 3), keepdim=True)
            if self.elementwise_affine:
                x = (
                        self.weight * (x - mean) / torch.sqrt(var + self.eps)
                        + self.bias
                )
            else:
                x = (x - mean) / torch.sqrt(var + self.eps)
        return x


class CumulativeLayerNorm(nn.LayerNorm):
    """Calculate Cumulative Layer Normalization.

       Arguments
       ---------
       dim : int
        Dimension that you want to normalize.
       elementwise_affine : True
        Learnable per-element affine parameters.

    Example
    -------
    >>> x = torch.randn(5, 10, 20)
    >>> CLN = CumulativeLayerNorm(10)
    >>> x_norm = CLN(x)
    """

    def __init__(self, dim, elementwise_affine=True, eps=1e-8):
        super(CumulativeLayerNorm, self).__init__(
            dim, elementwise_affine=elementwise_affine, eps=eps
        )

    def forward(self, x):
        """Returns the normalized tensor.

        Arguments
        ---------
        x : torch.Tensor
            Tensor size [N, C, K, S] or [N, C, L]
        """
        # x: N x C x K x S or N x C x L
        # N x K x S x C
        if x.dim() == 4:
            x = x.permute(0, 2, 3, 1).contiguous()
            # N x K x S x C == only channel norm
            x = super().forward(x)
            # N x C x K x S
            x = x.permute(0, 3, 1, 2).contiguous()
        if x.dim() == 3:
            x = torch.transpose(x, 1, 2)
            # N x L x C == only channel norm
            x = super().forward(x)
            # N x C x L
            x = torch.transpose(x, 1, 2)
        return x


def select_norm(norm, dim, shape, eps=1e-8):
    """Just a wrapper to select the normalization type.
    """

    if norm == "gln":
        return GlobalLayerNorm(dim, shape, elementwise_affine=True, eps=eps)
    if norm == "cln":
        return CumulativeLayerNorm(dim, elementwise_affine=True, eps=eps)
    if norm == "ln":
        return nn.GroupNorm(1, dim, eps=eps)
    else:
        return nn.BatchNorm1d(dim)


class Encoder(nn.Module):
    """Convolutional Encoder Layer.

    Arguments
    ---------
    kernel_size : int
        Length of filters.
    in_channels : int
        Number of  input channels.
    out_channels : int
        Number of output channels.

    Example
    -------
    >>> x = torch.randn(2, 1000)
    >>> encoder = Encoder(kernel_size=4, out_channels=64)
    >>> h = encoder(x)
    >>> h.shape
    torch.Size([2, 64, 499])
    """

    def __init__(self, kernel_size=2, out_channels=64, in_channels=1):
        super(Encoder, self).__init__()
        self.conv1d = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=kernel_size // 2,
            groups=1,
            bias=False,
        )
        self.in_channels = in_channels

    def forward(self, x):
        """Return the encoded output.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor with dimensionality [B, L].
        Return
        ------
        x : torch.Tensor
            Encoded tensor with dimensionality [B, N, T_out].

        where B = Batchsize
              L = Number of timepoints
              N = Number of filters
              T_out = Number of timepoints at the output of the encoder
        """
        # B x L -> B x 1 x L
        if self.in_channels == 1:
            x = torch.unsqueeze(x, dim=1)
        # B x 1 x L -> B x N x T_out
        x = self.conv1d(x)
        x = F.relu(x)

        return x


class Decoder(nn.ConvTranspose1d):
    """A decoder layer that consists of ConvTranspose1d.

    Arguments
    ---------
    kernel_size : int
        Length of filters.
    in_channels : int
        Number of  input channels.
    out_channels : int
        Number of output channels.


    Example
    ---------
    >>> x = torch.randn(2, 100, 1000)
    >>> decoder = Decoder(kernel_size=4, in_channels=100, out_channels=1)
    >>> h = decoder(x)
    >>> h.shape
    torch.Size([2, 1003])
    """

    def __init__(self, *args, **kwargs):
        super(Decoder, self).__init__(*args, **kwargs)

    def forward(self, x):
        """Return the decoded output.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor with dimensionality [B, N, L].
                where, B = Batchsize,
                       N = number of filters
                       L = time points
        """

        if x.dim() not in [2, 3]:
            raise RuntimeError(
                "{} accept 3/4D tensor as input".format(self.__name__)
            )
        x = super().forward(x if x.dim() == 3 else torch.unsqueeze(x, 1))

        if torch.squeeze(x).dim() == 1:
            x = torch.squeeze(x, dim=1)
        else:
            x = torch.squeeze(x)
        return x


class IdentityBlock:
    """This block is used when we want to have identity transformation within the Dual_path block.

    Example
    -------
    >>> x = torch.randn(10, 100)
    >>> IB = IdentityBlock()
    >>> xhat = IB(x)
    """

    def _init__(self, **kwargs):
        pass

    def __call__(self, x):
        return x


class FastTransformerBlock(nn.Module):
    """This block is used to implement fast transformer models with efficient attention.

    The implementations are taken from https://fast-transformers.github.io/

    Arguments
    ---------
    attention_type : str
        Specifies the type of attention.
        Check https://fast-transformers.github.io/  for details.
    out_channels : int
        Dimensionality of the representation.
    num_layers : int
        Number of layers.
    nhead : int
        Number of attention heads.
    d_ffn : int
        Dimensionality of positional feed-forward.
    dropout : float
        Dropout drop rate.
    activation : str
        Activation function.
    reformer_bucket_size : int
        bucket size for reformer.

    Example
    -------
    # >>> x = torch.randn(10, 100, 64)
    # >>> block = FastTransformerBlock('linear', 64)
    # >>> x = block(x)
    # >>> x.shape
    # torch.Size([10, 100, 64])
    """

    def __init__(
            self,
            attention_type,
            out_channels,
            num_layers=6,
            nhead=8,
            d_ffn=1024,
            dropout=0,
            activation="relu",
            reformer_bucket_size=32,
    ):
        super(FastTransformerBlock, self).__init__()
        from fast_transformers.builders import TransformerEncoderBuilder

        builder = TransformerEncoderBuilder.from_kwargs(
            attention_type=attention_type,
            n_layers=num_layers,
            n_heads=nhead,
            feed_forward_dimensions=d_ffn,
            query_dimensions=out_channels // nhead,
            value_dimensions=out_channels // nhead,
            dropout=dropout,
            attention_dropout=dropout,
            chunk_size=reformer_bucket_size,
        )
        self.mdl = builder.get()

        self.attention_type = attention_type
        self.reformer_bucket_size = reformer_bucket_size

    def forward(self, x):
        """Returns the transformed input.

        Arguments
        ---------
        x : torch.Tensor
            Tensor shaper [B, L, N].
            where, B = Batchsize,
                   N = number of filters
                   L = time points
        """
        if self.attention_type == "reformer":

            # pad zeros at the end
            pad_size = (self.reformer_bucket_size * 2) - (
                    x.shape[1] % (self.reformer_bucket_size * 2)
            )
            device = x.device
            x_padded = torch.cat(
                [x, torch.zeros(x.size(0), pad_size, x.size(-1)).to(device)],
                dim=1,
            )

            # apply the model
            x_padded = self.mdl(x_padded)

            # get rid of zeros at the end
            return x_padded[:, :-pad_size, :]
        else:
            return self.mdl(x)


class PyTorchPositionalEncoding(nn.Module):
    """Positional encoder for the pytorch transformer.

    Arguments
    ---------
    d_model : int
        Representation dimensionality.
    dropout : float
        Dropout drop prob.
    max_len : int
        Max sequence length.

    Example
    -------
    >>> x = torch.randn(10, 100, 64)
    >>> enc = PyTorchPositionalEncoding(64)
    >>> x = enc(x)
    """

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PyTorchPositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer("pe", pe)

    def forward(self, x):
        """Returns the encoded output.

        Arguments
        ---------
        x : torch.Tensor
            Tensor shape [B, L, N],
            where, B = Batchsize,
                   N = number of filters
                   L = time points
        """
        x = x + self.pe[: x.size(0), :]
        return self.dropout(x)


class PytorchTransformerBlock(nn.Module):
    """A wrapper that uses the pytorch transformer block.

    Arguments
    ---------
    out_channels : int
        Dimensionality of the representation.
    num_layers : int
        Number of layers.
    nhead : int
        Number of attention heads.
    d_ffn : int
        Dimensionality of positional feed forward.
    Dropout : float
        Dropout drop rate.
    activation : str
        Activation function.
    use_positional_encoding : bool
        If true we use a positional encoding.

    Example
    ---------
    >>> x = torch.randn(10, 100, 64)
    >>> block = PytorchTransformerBlock(64)
    >>> x = block(x)
    >>> x.shape
    torch.Size([10, 100, 64])
    """

    def __init__(
            self,
            out_channels,
            num_layers=6,
            nhead=8,
            d_ffn=2048,
            dropout=0.1,
            activation="relu",
            use_positional_encoding=True,
    ):
        super(PytorchTransformerBlock, self).__init__()

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=out_channels,
            nhead=nhead,
            dim_feedforward=d_ffn,
            dropout=dropout,
            activation=activation,
        )
        # cem :this encoder thing has a normalization component. we should look at that probably also.
        self.mdl = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        if use_positional_encoding:
            self.pos_encoder = PyTorchPositionalEncoding(out_channels)
        else:
            self.pos_encoder = None

    def forward(self, x):
        """Returns the transformed output.

        Arguments
        ---------
        x : torch.Tensor
            Tensor shape [B, L, N]
            where, B = Batchsize,
                   N = number of filters
                   L = time points

        """
        if self.pos_encoder is not None:
            x = self.pos_encoder(x)
        return self.mdl(x)


class SBTransformerBlock(nn.Module):
    """A wrapper for the SpeechBrain implementation of the transformer encoder.

    Arguments
    ---------
    num_layers : int
        Number of layers.
    d_model : int
        Dimensionality of the representation.
    nhead : int
        Number of attention heads.
    d_ffn : int
        Dimensionality of positional feed forward.
    input_shape : tuple
        Shape of input.
    kdim : int
        Dimension of the key (Optional).
    vdim : int
        Dimension of the value (Optional).
    dropout : float
        Dropout rate.
    activation : str
        Activation function.
    use_positional_encoding : bool
        If true we use a positional encoding.
    norm_before: bool
        Use normalization before transformations.

    Example
    ---------
    >>> x = torch.randn(10, 100, 64)
    >>> block = SBTransformerBlock(1, 64, 8)
    >>> x = block(x)
    >>> x.shape
    torch.Size([10, 100, 64])
    """

    def __init__(
            self,
            num_layers,
            d_model,
            nhead,
            d_ffn=2048,
            input_shape=None,
            kdim=None,
            vdim=None,
            dropout=0.1,
            activation="relu",
            use_positional_encoding=False,
            norm_before=False,
            attention_type="regularMHA",
    ):
        super(SBTransformerBlock, self).__init__()
        self.use_positional_encoding = use_positional_encoding

        if activation == "relu":
            activation = nn.ReLU
        elif activation == "gelu":
            activation = nn.GELU
        else:
            raise ValueError("unknown activation")

        self.mdl = TransformerEncoder(
            num_layers=num_layers,
            nhead=nhead,
            d_ffn=d_ffn,
            input_shape=input_shape,
            d_model=d_model,
            kdim=kdim,
            vdim=vdim,
            dropout=dropout,
            activation=activation,
            normalize_before=norm_before,
            attention_type=attention_type,
        )

        if use_positional_encoding:
            self.pos_enc = PositionalEncoding(input_size=d_model)

    def forward(self, x):
        """Returns the transformed output.

        Arguments
        ---------
        x : torch.Tensor
            Tensor shape [B, L, N],
            where, B = Batchsize,
                   L = time points
                   N = number of filters

        """
        if self.use_positional_encoding:
            pos_enc = self.pos_enc(x)
            return self.mdl(x + pos_enc)[0]
        else:
            return self.mdl(x)[0]


class SBRNNBlock(nn.Module):
    """RNNBlock for the dual path pipeline.

    Arguments
    ---------
    input_size : int
        Dimensionality of the input features.
    hidden_channels : int
        Dimensionality of the latent layer of the rnn.
    num_layers : int
        Number of the rnn layers.
    rnn_type : str
        Type of the the rnn cell.
    dropout : float
        Dropout rate
    bidirectional : bool
        If True, bidirectional.

    Example
    ---------
    >>> x = torch.randn(10, 100, 64)
    >>> rnn = SBRNNBlock(64, 100, 1, bidirectional=True)
    >>> x = rnn(x)
    >>> x.shape
    torch.Size([10, 100, 200])
    """

    def __init__(
            self,
            input_size,
            hidden_channels,
            num_layers,
            rnn_type="LSTM",
            dropout=0,
            bidirectional=True,
    ):
        super(SBRNNBlock, self).__init__()

        self.mdl = getattr(SBRNN, rnn_type)(
            hidden_channels,
            input_size=input_size,
            num_layers=num_layers,
            dropout=dropout,
            bidirectional=bidirectional,
        )

    def forward(self, x):
        """Returns the transformed output.

        Arguments
        ---------
        x : torch.Tensor
            [B, L, N]
            where, B = Batchsize,
                   N = number of filters
                   L = time points
        """

        return self.mdl(x)[0]


class DPTNetBlock(nn.Module):
    """The DPT Net block.

    Arguments
    ---------
    d_model : int
        Number of expected features in the input (required).
    nhead : int
        Number of heads in the multiheadattention models (required).
    dim_feedforward : int
        Dimension of the feedforward network model (default=2048).
    dropout : float
        Dropout value (default=0.1).
    activation : str
        Activation function of intermediate layer, relu or gelu (default=relu).

    Examples
    --------
        >>> encoder_layer = DPTNetBlock(d_model=512, nhead=8)
        >>> src = torch.rand(10, 100, 512)
        >>> out = encoder_layer(src)
        >>> out.shape
        torch.Size([10, 100, 512])
    """

    def __init__(
            self, d_model, nhead, dim_feedforward=256, dropout=0, activation="relu"
    ):
        from torch.nn.modules.activation import MultiheadAttention
        from torch.nn.modules.normalization import LayerNorm
        from torch.nn.modules.dropout import Dropout
        from torch.nn.modules.rnn import LSTM
        from torch.nn.modules.linear import Linear

        super(DPTNetBlock, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        # self.linear1 = Linear(d_model, dim_feedforward)
        self.rnn = LSTM(d_model, d_model * 2, 1, bidirectional=True)
        self.dropout = Dropout(dropout)
        # self.linear2 = Linear(dim_feedforward, d_model)
        self.linear2 = Linear(d_model * 2 * 2, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if "activation" not in state:
            state["activation"] = F.relu
        super(DPTNetBlock, self).__setstate__(state)

    def forward(self, src):
        """Pass the input through the encoder layer.

        Arguments
        ---------
        src : torch.Tensor
            Tensor shape [B, L, N]
            where, B = Batchsize,
                   N = number of filters
                   L = time points

        """
        src2 = self.self_attn(
            src, src, src, attn_mask=None, key_padding_mask=None
        )[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        # src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src2 = self.rnn(src)[0]
        src2 = self.activation(src2)
        src2 = self.dropout(src2)
        src2 = self.linear2(src2)
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


def _get_activation_fn(activation):
    """Just a wrapper to get the activation functions.
    """

    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu


class Dual_Computation_Block(nn.Module):
    """Computation block for dual-path processing.

    Arguments
    ---------
    intra_mdl : torch.nn.module
        Model to process within the chunks.
     inter_mdl : torch.nn.module
        Model to process across the chunks.
     out_channels : int
        Dimensionality of inter/intra model.
     norm : str
        Normalization type.
     skip_around_intra : bool
        Skip connection around the intra layer.
     linear_layer_after_inter_intra : bool
        Linear layer or not after inter or intra.

    Example
    ---------
        >>> intra_block = SBTransformerBlock(1, 64, 8)
        >>> inter_block = SBTransformerBlock(1, 64, 8)
        >>> dual_comp_block = Dual_Computation_Block(intra_block, inter_block, 64)
        >>> x = torch.randn(10, 64, 100, 10)
        >>> x = dual_comp_block(x)
        >>> x.shape
        torch.Size([10, 64, 100, 10])
    """

    def __init__(
            self,
            intra_mdl,
            inter_mdl,
            out_channels,
            norm="ln",
            skip_around_intra=True,
            linear_layer_after_inter_intra=True,
    ):
        super(Dual_Computation_Block, self).__init__()

        self.intra_mdl = intra_mdl
        self.inter_mdl = inter_mdl
        self.skip_around_intra = skip_around_intra
        self.linear_layer_after_inter_intra = linear_layer_after_inter_intra

        # Norm
        self.norm = norm
        if norm is not None:
            self.intra_norm = select_norm(norm, out_channels, 4)
            self.inter_norm = select_norm(norm, out_channels, 4)

        # Linear
        if linear_layer_after_inter_intra:
            if isinstance(intra_mdl, SBRNNBlock):
                self.intra_linear = Linear(
                    out_channels, input_size=2 * intra_mdl.mdl.rnn.hidden_size
                )
            else:
                self.intra_linear = Linear(
                    out_channels, input_size=out_channels
                )

            if isinstance(inter_mdl, SBRNNBlock):
                self.inter_linear = Linear(
                    out_channels, input_size=2 * intra_mdl.mdl.rnn.hidden_size
                )
            else:
                self.inter_linear = Linear(
                    out_channels, input_size=out_channels
                )

    def forward(self, x):
        """Returns the output tensor.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor of dimension [B, N, K, S].


        Return
        ---------
        out: torch.Tensor
            Output tensor of dimension [B, N, K, S].
            where, B = Batchsize,
               N = number of filters
               K = time points in each chunk
               S = the number of chunks
        """
        B, N, K, S = x.shape
        # intra RNN
        # [BS, K, N]
        intra = x.permute(0, 3, 2, 1).contiguous().view(B * S, K, N)
        # [BS, K, H]

        intra = self.intra_mdl(intra)

        # [BS, K, N]
        if self.linear_layer_after_inter_intra:
            intra = self.intra_linear(intra)

        # [B, S, K, N]
        intra = intra.view(B, S, K, N)
        # [B, N, K, S]
        intra = intra.permute(0, 3, 2, 1).contiguous()
        if self.norm is not None:
            intra = self.intra_norm(intra)

        # [B, N, K, S]
        if self.skip_around_intra:
            intra = intra + x

        # inter RNN
        # [BK, S, N]
        inter = intra.permute(0, 2, 3, 1).contiguous().view(B * K, S, N)
        # [BK, S, H]
        inter = self.inter_mdl(inter)

        # [BK, S, N]
        if self.linear_layer_after_inter_intra:
            inter = self.inter_linear(inter)

        # [B, K, S, N]
        inter = inter.view(B, K, S, N)
        # [B, N, K, S]
        inter = inter.permute(0, 3, 1, 2).contiguous()
        if self.norm is not None:
            inter = self.inter_norm(inter)
        # [B, N, K, S]
        out = inter + intra

        return out


class Dual_Path_Model(nn.Module):
    """The dual path model which is the basis for dualpathrnn, sepformer, dptnet.

    Arguments
    ---------
    in_channels : int
        Number of channels at the output of the encoder.
    out_channels : int
        Number of channels that would be inputted to the intra and inter blocks.
    intra_model : torch.nn.module
        Model to process within the chunks.
    inter_model : torch.nn.module
        model to process across the chunks,
    num_layers : int
        Number of layers of Dual Computation Block.
    norm : str
        Normalization type.
    K : int
        Chunk length.
    num_spks : int
        Number of sources (speakers).
    skip_around_intra : bool
        Skip connection around intra.
    linear_layer_after_inter_intra : bool
        Linear layer after inter and intra.
    use_global_pos_enc : bool
        Global positional encodings.
    max_length : int
        Maximum sequence length.

    Example
    ---------
    >>> intra_block = SBTransformerBlock(1, 64, 8)
    >>> inter_block = SBTransformerBlock(1, 64, 8)
    >>> dual_path_model = Dual_Path_Model(64, 64, intra_block, inter_block, num_spks=2)
    >>> x = torch.randn(10, 64, 2000)
    >>> x = dual_path_model(x)
    >>> x.shape
    torch.Size([2, 10, 64, 2000])
    """

    def __init__(
            self,
            in_channels,
            out_channels,
            intra_model,
            inter_model,
            num_layers=1,
            norm="ln",
            K=200,
            num_spks=2,
            skip_around_intra=True,
            linear_layer_after_inter_intra=True,
            use_global_pos_enc=False,
            max_length=20000,
            use_noise_ground=False,
            use_noise_ground_only=False
    ):
        super(Dual_Path_Model, self).__init__()
        self.K = K
        self.num_spks = num_spks
        self.num_layers = num_layers
        self.norm = select_norm(norm, in_channels, 3)
        self.conv1d = nn.Conv1d(in_channels, out_channels, 1, bias=False)
        self.use_global_pos_enc = use_global_pos_enc

        if self.use_global_pos_enc:
            self.pos_enc = PositionalEncoding(max_length)

        self.dual_mdl = nn.ModuleList([])
        for i in range(num_layers):
            self.dual_mdl.append(
                copy.deepcopy(
                    Dual_Computation_Block(
                        intra_model,
                        inter_model,
                        out_channels,
                        norm,
                        skip_around_intra=skip_around_intra,
                        linear_layer_after_inter_intra=linear_layer_after_inter_intra,
                    )
                )
            )

        self.use_noise_ground = use_noise_ground
        self.use_noise_ground_only = use_noise_ground_only

        self.conv2d = nn.Conv2d(
            out_channels, out_channels * (num_spks + 1), kernel_size=1
        ) if self.use_noise_ground and not self.use_noise_ground_only else nn.Conv2d(
            out_channels, out_channels * num_spks, kernel_size=1
        )

        self.end_conv1x1 = nn.Conv1d(out_channels, in_channels, 1, bias=False)
        self.prelu = nn.PReLU()
        self.activation = nn.ReLU()
        # gated output layer
        self.output = nn.Sequential(
            nn.Conv1d(out_channels, out_channels, 1), nn.Tanh()
        )
        self.output_gate = nn.Sequential(
            nn.Conv1d(out_channels, out_channels, 1), nn.Sigmoid()
        )

    def forward(self, x):
        """Returns the output tensor.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor of dimension [B, N, L].

        Returns
        -------
        out : torch.Tensor
            Output tensor of dimension [spks, B, N, L]
            where, spks = Number of speakers
               B = Batchsize,
               N = number of filters
               L = the number of time points
        """

        # before each line we indicate the shape after executing the line

        # [B, N, L]
        x = self.norm(x)

        # [B, N, L]
        x = self.conv1d(x)
        if self.use_global_pos_enc:
            x = self.pos_enc(x.transpose(1, -1)).transpose(1, -1) + x * (
                    x.size(1) ** 0.5
            )

        # [B, N, K, S]
        x, gap = self._Segmentation(x, self.K)

        # [B, N, K, S]
        for i in range(self.num_layers):
            x = self.dual_mdl[i](x)
        x = self.prelu(x)

        # [B, N*spks, K, S]
        x = self.conv2d(x)
        B, _, K, S = x.shape

        num_spks = self.num_spks + 1 if self.use_noise_ground and not self.use_noise_ground_only else self.num_spks

        # [B*spks, N, K, S]
        x = x.view(B * num_spks, -1, K, S)

        # [B*spks, N, L]
        x = self._over_add(x, gap)
        x = self.output(x) * self.output_gate(x)

        # [B*spks, N, L]
        x = self.end_conv1x1(x)

        # [B, spks, N, L]
        _, N, L = x.shape
        x = x.view(B, num_spks, N, L)
        x = self.activation(x)

        # [spks, B, N, L]
        x = x.transpose(0, 1)

        return x

    def _padding(self, input, K):
        """Padding the audio times.

        Arguments
        ---------
        K : int
            Chunks of length.
        P : int
            Hop size.
        input : torch.Tensor
            Tensor of size [B, N, L].
            where, B = Batchsize,
                   N = number of filters
                   L = time points
        """
        B, N, L = input.shape
        P = K // 2
        gap = K - (P + L % K) % K
        if gap > 0:
            pad = (
                torch.Tensor(torch.zeros(B, N, gap))
                .type(input.dtype)
                .to(input.device)
            )
            input = torch.cat([input, pad], dim=2)

        _pad = (
            torch.Tensor(torch.zeros(B, N, P))
            .type(input.dtype)
            .to(input.device)
        )
        input = torch.cat([_pad, input, _pad], dim=2)

        return input, gap

    def _Segmentation(self, input, K):
        """The segmentation stage splits

        Arguments
        ---------
        K : int
            Length of the chunks.
        input : torch.Tensor
            Tensor with dim [B, N, L].

        Return
        -------
        output : torch.tensor
            Tensor with dim [B, N, K, S].
            where, B = Batchsize,
               N = number of filters
               K = time points in each chunk
               S = the number of chunks
               L = the number of time points
        """
        B, N, L = input.shape
        P = K // 2
        input, gap = self._padding(input, K)
        # [B, N, K, S]
        input1 = input[:, :, :-P].contiguous().view(B, N, -1, K)
        input2 = input[:, :, P:].contiguous().view(B, N, -1, K)
        input = (
            torch.cat([input1, input2], dim=3).view(B, N, -1, K).transpose(2, 3)
        )

        return input.contiguous(), gap

    def _over_add(self, input, gap):
        """Merge the sequence with the overlap-and-add method.

        Arguments
        ---------
        input : torch.tensor
            Tensor with dim [B, N, K, S].
        gap : int
            Padding length.

        Return
        -------
        output : torch.tensor
            Tensor with dim [B, N, L].
            where, B = Batchsize,
               N = number of filters
               K = time points in each chunk
               S = the number of chunks
               L = the number of time points

        """
        B, N, K, S = input.shape
        P = K // 2
        # [B, N, S, K]
        input = input.transpose(2, 3).contiguous().view(B, N, -1, K * 2)

        input1 = input[:, :, :, :K].contiguous().view(B, N, -1)[:, :, P:]
        input2 = input[:, :, :, K:].contiguous().view(B, N, -1)[:, :, :-P]
        input = input1 + input2
        # [B, N, L]
        if gap > 0:
            input = input[:, :, :-gap]

        return input


class SepformerWrapper(nn.Module):
    """The wrapper for the sepformer model which combines the Encoder, Masknet and the decoder
    https://arxiv.org/abs/2010.13154

    Arguments
    ---------

    encoder_kernel_size: int,
        The kernel size used in the encoder
    encoder_in_nchannels: int,
        The number of channels of the input audio
    encoder_out_nchannels: int,
        The number of filters used in the encoder.
        Also, number of channels that would be inputted to the intra and inter blocks.
    masknet_chunksize: int,
        The chunk length that is to be processed by the intra blocks
    masknet_numlayers: int,
        The number of layers of combination of inter and intra blocks
    masknet_norm: str,
        The normalization type to be used in the masknet
        Should be one of 'ln' -- layernorm, 'gln' -- globallayernorm
                         'cln' -- cumulative layernorm, 'bn' -- batchnorm
                         -- see the select_norm function above for more details
    masknet_useextralinearlayer: bool,
        Whether or not to use a linear layer at the output of intra and inter blocks
    masknet_extraskipconnection: bool,
        This introduces extra skip connections around the intra block
    masknet_numspks: int,
        This determines the number of speakers to estimate
    intra_numlayers: int,
        This determines the number of layers in the intra block
    inter_numlayers: int,
        This determines the number of layers in the inter block
    intra_nhead: int,
        This determines the number of parallel attention heads in the intra block
    inter_nhead: int,
        This determines the number of parallel attention heads in the inter block
    intra_dffn: int,
        The number of dimensions in the positional feedforward model in the inter block
    inter_dffn: int,
        The number of dimensions in the positional feedforward model in the intra block
    intra_use_positional: bool,
        Whether or not to use positional encodings in the intra block
    inter_use_positional: bool,
        Whether or not to use positional encodings in the inter block
    intra_norm_before: bool
        Whether or not we use normalization before the transformations in the intra block
    inter_norm_before: bool
        Whether or not we use normalization before the transformations in the inter block

    Example
    -----
    >>> model = SepformerWrapper()
    >>> inp = torch.rand(1, 160)
    >>> result = model.forward(inp)
    >>> result.shape
    torch.Size([1, 160, 2])
    """

    def __init__(
            self,
            encoder_kernel_size=16,
            encoder_in_nchannels=1,
            encoder_out_nchannels=256,
            masknet_chunksize=250,
            masknet_numlayers=2,
            masknet_norm="ln",
            masknet_useextralinearlayer=False,
            masknet_extraskipconnection=True,
            masknet_numspks=2,
            intra_numlayers=8,
            inter_numlayers=8,
            intra_nhead=8,
            inter_nhead=8,
            intra_dffn=1024,
            inter_dffn=1024,
            intra_use_positional=True,
            inter_use_positional=True,
            intra_norm_before=True,
            inter_norm_before=True,
    ):

        super(SepformerWrapper, self).__init__()
        self.encoder = Encoder(
            kernel_size=encoder_kernel_size,
            out_channels=encoder_out_nchannels,
            in_channels=encoder_in_nchannels,
        )
        intra_model = SBTransformerBlock(
            num_layers=intra_numlayers,
            d_model=encoder_out_nchannels,
            nhead=intra_nhead,
            d_ffn=intra_dffn,
            use_positional_encoding=intra_use_positional,
            norm_before=intra_norm_before,
        )

        inter_model = SBTransformerBlock(
            num_layers=inter_numlayers,
            d_model=encoder_out_nchannels,
            nhead=inter_nhead,
            d_ffn=inter_dffn,
            use_positional_encoding=inter_use_positional,
            norm_before=inter_norm_before,
        )

        self.masknet = Dual_Path_Model(
            in_channels=encoder_out_nchannels,
            out_channels=encoder_out_nchannels,
            intra_model=intra_model,
            inter_model=inter_model,
            num_layers=masknet_numlayers,
            norm=masknet_norm,
            K=masknet_chunksize,
            num_spks=masknet_numspks,
            skip_around_intra=masknet_extraskipconnection,
            linear_layer_after_inter_intra=masknet_useextralinearlayer,
        )
        self.decoder = Decoder(
            in_channels=encoder_out_nchannels,
            out_channels=encoder_in_nchannels,
            kernel_size=encoder_kernel_size,
            stride=encoder_kernel_size // 2,
            bias=False,
        )
        self.num_spks = masknet_numspks

        # reinitialize the parameters
        for module in [self.encoder, self.masknet, self.decoder]:
            self.reset_layer_recursively(module)

    def reset_layer_recursively(self, layer):
        """Reinitializes the parameters of the network"""
        if hasattr(layer, "reset_parameters"):
            layer.reset_parameters()
        for child_layer in layer.modules():
            if layer != child_layer:
                self.reset_layer_recursively(child_layer)

    def forward(self, mix):
        """ Processes the input tensor x and returns an output tensor."""
        mix_w = self.encoder(mix)
        est_mask = self.masknet(mix_w)
        mix_w = torch.stack([mix_w] * self.num_spks)
        sep_h = mix_w * est_mask

        # Decoding
        est_source = torch.cat(
            [
                self.decoder(sep_h[i]).unsqueeze(-1)
                for i in range(self.num_spks)
            ],
            dim=-1,
        )

        # T changed after conv1d in encoder, fix it here
        T_origin = mix.size(1)
        T_est = est_source.size(1)
        if T_origin > T_est:
            est_source = F.pad(est_source, (0, 0, 0, T_origin - T_est))
        else:
            est_source = est_source[:, :T_origin, :]

        return est_source


class SBConformerEncoderBlock(nn.Module):
    """A wrapper for the SpeechBrain implementation of the ConformerEncoder.

    Arguments
    ---------
    num_layers : int
        Number of layers.
    d_model : int
        Dimensionality of the representation.
    nhead : int
        Number of attention heads.
    d_ffn : int
        Dimensionality of positional feed forward.
    input_shape : tuple
        Shape of input.
    kdim : int
        Dimension of the key (Optional).
    vdim : int
        Dimension of the value (Optional).
    dropout : float
        Dropout rate.
    activation : str
        Activation function.
    kernel_size: int
        Kernel size in the conformer encoder
    bias: bool
        Use bias or not in the convolution part of conformer encoder
    use_positional_encoding : bool
        If true we use a positional encoding.


    Example
    ---------
    >>> x = torch.randn(10, 100, 64)
    >>> block = SBConformerEncoderBlock(1, 64, 8)
    >>> from speechbrain.lobes.models.transformer.Transformer import PositionalEncoding
    >>> pos_enc = PositionalEncoding(64)
    >>> pos_embs = pos_enc(torch.ones(1, 199, 64))
    >>> x = block(x)
    >>> x.shape
    torch.Size([10, 100, 64])
    """

    def __init__(
            self,
            num_layers,
            d_model,
            nhead,
            d_ffn=2048,
            input_shape=None,
            kdim=None,
            vdim=None,
            dropout=0.1,
            activation="swish",
            kernel_size=31,
            bias=True,
            use_positional_encoding=True,
            attention_type="RelPosMHAXL",
    ):
        super(SBConformerEncoderBlock, self).__init__()
        self.use_positional_encoding = use_positional_encoding
        self.attention_type = attention_type

        if activation == "relu":
            activation = nn.ReLU
        elif activation == "gelu":
            activation = nn.GELU
        elif activation == "swish":
            activation = Swish
        else:
            raise ValueError("unknown activation")

        self.mdl = ConformerEncoder(
            num_layers=num_layers,
            nhead=nhead,
            d_ffn=d_ffn,
            d_model=d_model,
            kdim=kdim,
            vdim=vdim,
            dropout=dropout,
            activation=activation,
            kernel_size=kernel_size,
            bias=bias,
            attention_type=attention_type,
        )

        if self.attention_type == "RelPosMHAXL":
            # for RelPosMHAXL, we need the positional encoding (not optional)
            self.pos_enc = PositionalEncoding(input_size=d_model)
        elif self.attention_type == "regularMHA":
            if self.use_positional_encoding:
                self.pos_enc = PositionalEncoding(input_size=d_model)
        else:
            raise ValueError("Unsupported attention type")

    def forward(self, x):
        """Returns the transformed output.

        Arguments
        ---------
        x : torch.Tensor
            Tensor shape [B, L, N],
            where, B = Batchsize,
                   L = time points
                   N = number of filters

        """
        if self.attention_type == "RelPosMHAXL":
            pos_enc = self.pos_enc(
                torch.ones(
                    x.shape[0], x.shape[1] * 2 - 1, x.shape[2], device=x.device
                )
            )
            return self.mdl(x, pos_embs=pos_enc)[0]
        elif self.attention_type == "regularMHA":
            if self.use_positional_encoding:
                pos_embs = self.pos_enc(x)
                return self.mdl(x + pos_embs)[0]
            else:
                return self.mdl(x)[0]
        else:
            raise ValueError("Unsupported attention type")

speechbrain/lobes/models/dual_path.py
TzuchengChang/NASS
import torch
import torch.nn as nn
import numpy as np


class Normalize(nn.Module):

    def __init__(self, power=2):
        super(Normalize, self).__init__()
        self.power = power

    def forward(self, x):
        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)
        out = x.div(norm + 1e-7)
        return out


class PatchSampleF(nn.Module):
    def __init__(self, use_mlp=False, init_type='normal', init_gain=0.02, nc=256, gpu_ids=[], num_smp=2):
        super(PatchSampleF, self).__init__()
        self.l2norm = Normalize(2)
        self.use_mlp = use_mlp
        self.nc = nc
        self.init_type = init_type
        self.init_gain = init_gain
        self.gpu_ids = gpu_ids
        self.num_smp = num_smp
        conv = nn.Sequential(*[nn.Conv2d(1, 1, 3, 1), nn.ReLU(), nn.Conv2d(1, 1, 3, 1)])
        mlp = nn.Sequential(*[nn.Linear(1, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)])
        if len(self.gpu_ids) > 0:
            conv.cuda()
            mlp.cuda()
        setattr(self, 'conv', conv)
        setattr(self, 'mlp', mlp)

    def forward(self, feats, num_patches=64, patch_ids=None):
        return_ids = []
        return_feats = []
        for feat_id, feat in enumerate(feats):
            if self.use_mlp:
                conv = getattr(self, 'conv')
                feat = conv(feat)
            B, H, W = feat.shape[0], feat.shape[2], feat.shape[3]
            feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2)
            if num_patches > 0:
                if patch_ids is not None:
                    patch_id = patch_ids[feat_id]
                else:
                    patch_id = np.random.permutation(feat_reshape.shape[1])
                    patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))]
                patch_id = torch.tensor(patch_id, dtype=torch.long, device=feat.device)
                x_sample = feat_reshape[:, patch_id, :].flatten(0, 1)
            else:
                x_sample = feat_reshape
                patch_id = []
            if self.use_mlp:
                mlp = getattr(self, 'mlp')
                x_sample = mlp(x_sample)
            return_ids.append(patch_id)
            x_sample = self.l2norm(x_sample)
            if num_patches == 0:
                x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W])
            return_feats.append(x_sample)
        return return_feats, return_ids

speechbrain/lobes/models/networks.py
TzuchengChang/NASS
"""
Losses for training neural networks.

Authors
 * Mirco Ravanelli 2020
 * Samuele Cornell 2020
 * Hwidong Na 2020
 * Yan Gao 2020
 * Titouan Parcollet 2020
"""

import math
import torch
import logging
import functools
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from itertools import permutations
from speechbrain.dataio.dataio import length_to_mask
from speechbrain.decoders.ctc import filter_ctc_output


logger = logging.getLogger(__name__)


def transducer_loss(
    logits,
    targets,
    input_lens,
    target_lens,
    blank_index,
    reduction="mean",
    use_torchaudio=True,
):
    """Transducer loss, see `speechbrain/nnet/loss/transducer_loss.py`.

    Arguments
    ---------
    logits : torch.Tensor
        Predicted tensor, of shape [batch, maxT, maxU, num_labels].
    targets : torch.Tensor
        Target tensor, without any blanks, of shape [batch, target_len].
    input_lens : torch.Tensor
        Length of each utterance.
    target_lens : torch.Tensor
        Length of each target sequence.
    blank_index : int
        The location of the blank symbol among the label indices.
    reduction : str
        Specifies the reduction to apply to the output: 'mean' | 'batchmean' | 'sum'.
    use_torchaudio: bool
        If True, use Transducer loss implementation from torchaudio, otherwise,
        use Speechbrain Numba implementation.
    """
    input_lens = (input_lens * logits.shape[1]).round().int()
    target_lens = (target_lens * targets.shape[1]).round().int()

    if use_torchaudio:
        try:
            from torchaudio.functional import rnnt_loss
        except ImportError:
            err_msg = "The dependency torchaudio >= 0.10.0 is needed to use Transducer Loss\n"
            err_msg += "Cannot import torchaudio.functional.rnnt_loss.\n"
            err_msg += "To use it, please install torchaudio >= 0.10.0\n"
            err_msg += "==================\n"
            err_msg += "Otherwise, you can use our numba implementation, set `use_torchaudio=False`.\n"
            raise ImportError(err_msg)

        return rnnt_loss(
            logits,
            targets.int(),
            input_lens,
            target_lens,
            blank=blank_index,
            reduction=reduction,
        )
    else:
        from speechbrain.nnet.loss.transducer_loss import Transducer

        # Transducer.apply function take log_probs tensor.
        log_probs = logits.log_softmax(-1)
        return Transducer.apply(
            log_probs, targets, input_lens, target_lens, blank_index, reduction,
        )


class PitWrapper(nn.Module):
    """
    Permutation Invariant Wrapper to allow Permutation Invariant Training
    (PIT) with existing losses.

    Permutation invariance is calculated over the sources/classes axis which is
    assumed to be the rightmost dimension: predictions and targets tensors are
    assumed to have shape [batch, ..., channels, sources].

    Arguments
    ---------
    base_loss : function
        Base loss function, e.g. torch.nn.MSELoss. It is assumed that it takes
        two arguments:
        predictions and targets and no reduction is performed.
        (if a pytorch loss is used, the user must specify reduction="none").

    Returns
    ---------
    pit_loss : torch.nn.Module
        Torch module supporting forward method for PIT.

    Example
    -------
    >>> pit_mse = PitWrapper(nn.MSELoss(reduction="none"))
    >>> targets = torch.rand((2, 32, 4))
    >>> p = (3, 0, 2, 1)
    >>> predictions = targets[..., p]
    >>> loss, opt_p = pit_mse(predictions, targets)
    >>> loss
    tensor([0., 0.])
    """

    def __init__(self, base_loss):
        super(PitWrapper, self).__init__()
        self.base_loss = base_loss

    def _fast_pit(self, loss_mat):
        """
        Arguments
        ----------
        loss_mat : torch.Tensor
            Tensor of shape [sources, source] containing loss values for each
            possible permutation of predictions.

        Returns
        -------
        loss : torch.Tensor
            Permutation invariant loss for the current batch, tensor of shape [1]

        assigned_perm : tuple
            Indexes for optimal permutation of the input over sources which
            minimizes the loss.
        """

        loss = None
        assigned_perm = None
        for p in permutations(range(loss_mat.shape[0])):
            c_loss = loss_mat[range(loss_mat.shape[0]), p].mean()
            if loss is None or loss > c_loss:
                loss = c_loss
                assigned_perm = p
        return loss, assigned_perm

    def _opt_perm_loss(self, pred, target):
        """
        Arguments
        ---------
        pred : torch.Tensor
            Network prediction for the current example, tensor of
            shape [..., sources].
        target : torch.Tensor
            Target for the current example, tensor of shape [..., sources].

        Returns
        -------
        loss : torch.Tensor
            Permutation invariant loss for the current example, tensor of shape [1]

        assigned_perm : tuple
            Indexes for optimal permutation of the input over sources which
            minimizes the loss.

        """

        n_sources = pred.size(-1)

        pred = pred.unsqueeze(-2).repeat(
            *[1 for x in range(len(pred.shape) - 1)], n_sources, 1
        )
        target = target.unsqueeze(-1).repeat(
            1, *[1 for x in range(len(target.shape) - 1)], n_sources
        )

        loss_mat = self.base_loss(pred, target)
        assert (
            len(loss_mat.shape) >= 2
        ), "Base loss should not perform any reduction operation"
        mean_over = [x for x in range(len(loss_mat.shape))]
        loss_mat = loss_mat.mean(dim=mean_over[:-2])

        return self._fast_pit(loss_mat)

    def reorder_tensor(self, tensor, p):
        """
        Arguments
        ---------
        tensor : torch.Tensor
            Tensor to reorder given the optimal permutation, of shape
            [batch, ..., sources].
        p : list of tuples
            List of optimal permutations, e.g. for batch=2 and n_sources=3
            [(0, 1, 2), (0, 2, 1].

        Returns
        -------
        reordered : torch.Tensor
            Reordered tensor given permutation p.
        """

        reordered = torch.zeros_like(tensor, device=tensor.device)
        for b in range(tensor.shape[0]):
            reordered[b] = tensor[b][..., p[b]].clone()
        return reordered

    def forward(self, preds, targets):
        """
            Arguments
            ---------
            preds : torch.Tensor
                Network predictions tensor, of shape
                [batch, channels, ..., sources].
            targets : torch.Tensor
                Target tensor, of shape [batch, channels, ..., sources].

            Returns
            -------
            loss : torch.Tensor
                Permutation invariant loss for current examples, tensor of
                shape [batch]

            perms : list
                List of indexes for optimal permutation of the inputs over
                sources.
                e.g., [(0, 1, 2), (2, 1, 0)] for three sources and 2 examples
                per batch.
        """
        losses = []
        perms = []
        for pred, label in zip(preds, targets):
            loss, p = self._opt_perm_loss(pred, label)
            perms.append(p)
            losses.append(loss)
        loss = torch.stack(losses)
        return loss, perms


def ctc_loss(
    log_probs, targets, input_lens, target_lens, blank_index, reduction="mean"
):
    """CTC loss.

    Arguments
    ---------
    predictions : torch.Tensor
        Predicted tensor, of shape [batch, time, chars].
    targets : torch.Tensor
        Target tensor, without any blanks, of shape [batch, target_len]
    input_lens : torch.Tensor
        Length of each utterance.
    target_lens : torch.Tensor
        Length of each target sequence.
    blank_index : int
        The location of the blank symbol among the character indexes.
    reduction : str
        What reduction to apply to the output. 'mean', 'sum', 'batch',
        'batchmean', 'none'.
        See pytorch for 'mean', 'sum', 'none'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.
    """
    input_lens = (input_lens * log_probs.shape[1]).round().int()
    target_lens = (target_lens * targets.shape[1]).round().int()
    log_probs = log_probs.transpose(0, 1)

    if reduction == "batchmean":
        reduction_loss = "sum"
    elif reduction == "batch":
        reduction_loss = "none"
    else:
        reduction_loss = reduction
    loss = torch.nn.functional.ctc_loss(
        log_probs,
        targets,
        input_lens,
        target_lens,
        blank_index,
        zero_infinity=True,
        reduction=reduction_loss,
    )

    if reduction == "batchmean":
        return loss / targets.shape[0]
    elif reduction == "batch":
        N = loss.size(0)
        return loss.view(N, -1).sum(1) / target_lens.view(N, -1).sum(1)
    else:
        return loss


def l1_loss(
    predictions, targets, length=None, allowed_len_diff=3, reduction="mean"
):
    """Compute the true l1 loss, accounting for length differences.

    Arguments
    ---------
    predictions : torch.Tensor
        Predicted tensor, of shape ``[batch, time, *]``.
    targets : torch.Tensor
        Target tensor with the same size as predicted tensor.
    length : torch.Tensor
        Length of each utterance for computing true error with a mask.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction : str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> probs = torch.tensor([[0.9, 0.1, 0.1, 0.9]])
    >>> l1_loss(probs, torch.tensor([[1., 0., 0., 1.]]))
    tensor(0.1000)
    """
    predictions, targets = truncate(predictions, targets, allowed_len_diff)
    loss = functools.partial(torch.nn.functional.l1_loss, reduction="none")
    return compute_masked_loss(
        loss, predictions, targets, length, reduction=reduction
    )


def mse_loss(
    predictions, targets, length=None, allowed_len_diff=3, reduction="mean"
):
    """Compute the true mean squared error, accounting for length differences.

    Arguments
    ---------
    predictions : torch.Tensor
        Predicted tensor, of shape ``[batch, time, *]``.
    targets : torch.Tensor
        Target tensor with the same size as predicted tensor.
    length : torch.Tensor
        Length of each utterance for computing true error with a mask.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction : str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> probs = torch.tensor([[0.9, 0.1, 0.1, 0.9]])
    >>> mse_loss(probs, torch.tensor([[1., 0., 0., 1.]]))
    tensor(0.0100)
    """
    predictions, targets = truncate(predictions, targets, allowed_len_diff)
    loss = functools.partial(torch.nn.functional.mse_loss, reduction="none")
    return compute_masked_loss(
        loss, predictions, targets, length, reduction=reduction
    )


def classification_error(
    probabilities, targets, length=None, allowed_len_diff=3, reduction="mean"
):
    """Computes the classification error at frame or batch level.

    Arguments
    ---------
    probabilities : torch.Tensor
        The posterior probabilities of shape
        [batch, prob] or [batch, frames, prob]
    targets : torch.Tensor
        The targets, of shape [batch] or [batch, frames]
    length : torch.Tensor
        Length of each utterance, if frame-level loss is desired.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction : str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> probs = torch.tensor([[[0.9, 0.1], [0.1, 0.9]]])
    >>> classification_error(probs, torch.tensor([1, 1]))
    tensor(0.5000)
    """
    if len(probabilities.shape) == 3 and len(targets.shape) == 2:
        probabilities, targets = truncate(
            probabilities, targets, allowed_len_diff
        )

    def error(predictions, targets):
        """Computes the classification error."""
        predictions = torch.argmax(probabilities, dim=-1)
        return (predictions != targets).float()

    return compute_masked_loss(
        error, probabilities, targets.long(), length, reduction=reduction
    )


def nll_loss(
    log_probabilities,
    targets,
    length=None,
    label_smoothing=0.0,
    allowed_len_diff=3,
    reduction="mean",
):
    """Computes negative log likelihood loss.

    Arguments
    ---------
    log_probabilities : torch.Tensor
        The probabilities after log has been applied.
        Format is [batch, log_p] or [batch, frames, log_p].
    targets : torch.Tensor
        The targets, of shape [batch] or [batch, frames].
    length : torch.Tensor
        Length of each utterance, if frame-level loss is desired.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction : str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> probs = torch.tensor([[0.9, 0.1], [0.1, 0.9]])
    >>> nll_loss(torch.log(probs), torch.tensor([1, 1]))
    tensor(1.2040)
    """
    if len(log_probabilities.shape) == 3:
        log_probabilities, targets = truncate(
            log_probabilities, targets, allowed_len_diff
        )
        log_probabilities = log_probabilities.transpose(1, -1)

    # Pass the loss function but apply reduction="none" first
    loss = functools.partial(torch.nn.functional.nll_loss, reduction="none")
    return compute_masked_loss(
        loss,
        log_probabilities,
        targets.long(),
        length,
        label_smoothing=label_smoothing,
        reduction=reduction,
    )


def bce_loss(
    inputs,
    targets,
    length=None,
    weight=None,
    pos_weight=None,
    reduction="mean",
    allowed_len_diff=3,
    label_smoothing=0.0,
):
    """Computes binary cross-entropy (BCE) loss. It also applies the sigmoid
    function directly (this improves the numerical stability).

    Arguments
    ---------
    inputs : torch.Tensor
        The output before applying the final softmax
        Format is [batch[, 1]?] or [batch, frames[, 1]?].
        (Works with or without a singleton dimension at the end).
    targets : torch.Tensor
        The targets, of shape [batch] or [batch, frames].
    length : torch.Tensor
        Length of each utterance, if frame-level loss is desired.
    weight : torch.Tensor
        A manual rescaling weight if provided it’s repeated to match input
        tensor shape.
    pos_weight : torch.Tensor
        A weight of positive examples. Must be a vector with length equal to
        the number of classes.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction: str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> inputs = torch.tensor([10.0, -6.0])
    >>> targets = torch.tensor([1, 0])
    >>> bce_loss(inputs, targets)
    tensor(0.0013)
    """
    # Squeeze singleton dimension so inputs + targets match
    if len(inputs.shape) == len(targets.shape) + 1:
        inputs = inputs.squeeze(-1)

    # Make sure tensor lengths match
    if len(inputs.shape) >= 2:
        inputs, targets = truncate(inputs, targets, allowed_len_diff)
    elif length is not None:
        raise ValueError("length can be passed only for >= 2D inputs.")

    # Pass the loss function but apply reduction="none" first
    loss = functools.partial(
        torch.nn.functional.binary_cross_entropy_with_logits,
        weight=weight,
        pos_weight=pos_weight,
        reduction="none",
    )
    return compute_masked_loss(
        loss,
        inputs,
        targets.float(),
        length,
        label_smoothing=label_smoothing,
        reduction=reduction,
    )


def kldiv_loss(
    log_probabilities,
    targets,
    length=None,
    label_smoothing=0.0,
    allowed_len_diff=3,
    pad_idx=0,
    reduction="mean",
):
    """Computes the KL-divergence error at the batch level.
    This loss applies label smoothing directly to the targets

    Arguments
    ---------
    probabilities : torch.Tensor
        The posterior probabilities of shape
        [batch, prob] or [batch, frames, prob].
    targets : torch.Tensor
        The targets, of shape [batch] or [batch, frames].
    length : torch.Tensor
        Length of each utterance, if frame-level loss is desired.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    reduction : str
        Options are 'mean', 'batch', 'batchmean', 'sum'.
        See pytorch for 'mean', 'sum'. The 'batch' option returns
        one loss per item in the batch, 'batchmean' returns sum / batch size.

    Example
    -------
    >>> probs = torch.tensor([[0.9, 0.1], [0.1, 0.9]])
    >>> kldiv_loss(torch.log(probs), torch.tensor([1, 1]))
    tensor(1.2040)
    """
    if label_smoothing > 0:
        if log_probabilities.dim() == 2:
            log_probabilities = log_probabilities.unsqueeze(1)

        bz, time, n_class = log_probabilities.shape
        targets = targets.long().detach()

        confidence = 1 - label_smoothing

        log_probabilities = log_probabilities.view(-1, n_class)
        targets = targets.view(-1)
        with torch.no_grad():
            true_distribution = log_probabilities.clone()
            true_distribution.fill_(label_smoothing / (n_class - 1))
            ignore = targets == pad_idx
            targets = targets.masked_fill(ignore, 0)
            true_distribution.scatter_(1, targets.unsqueeze(1), confidence)

        loss = torch.nn.functional.kl_div(
            log_probabilities, true_distribution, reduction="none"
        )
        loss = loss.masked_fill(ignore.unsqueeze(1), 0)

        # return loss according to reduction specified
        if reduction == "mean":
            return loss.sum().mean()
        elif reduction == "batchmean":
            return loss.sum() / bz
        elif reduction == "batch":
            return loss.view(bz, -1).sum(1) / length
        elif reduction == "sum":
            return loss.sum()
        else:
            return loss
    else:
        return nll_loss(log_probabilities, targets, length, reduction=reduction)


def truncate(predictions, targets, allowed_len_diff=3):
    """Ensure that predictions and targets are the same length.

    Arguments
    ---------
    predictions : torch.Tensor
        First tensor for checking length.
    targets : torch.Tensor
        Second tensor for checking length.
    allowed_len_diff : int
        Length difference that will be tolerated before raising an exception.
    """
    len_diff = predictions.shape[1] - targets.shape[1]
    if len_diff == 0:
        return predictions, targets
    elif abs(len_diff) > allowed_len_diff:
        raise ValueError(
            "Predictions and targets should be same length, but got %s and "
            "%s respectively." % (predictions.shape[1], targets.shape[1])
        )
    elif len_diff < 0:
        return predictions, targets[:, : predictions.shape[1]]
    else:
        return predictions[:, : targets.shape[1]], targets


def compute_masked_loss(
    loss_fn,
    predictions,
    targets,
    length=None,
    label_smoothing=0.0,
    reduction="mean",
):
    """Compute the true average loss of a set of waveforms of unequal length.

    Arguments
    ---------
    loss_fn : function
        A function for computing the loss taking just predictions and targets.
        Should return all the losses, not a reduction (e.g. reduction="none").
    predictions : torch.Tensor
        First argument to loss function.
    targets : torch.Tensor
        Second argument to loss function.
    length : torch.Tensor
        Length of each utterance to compute mask. If None, global average is
        computed and returned.
    label_smoothing: float
        The proportion of label smoothing. Should only be used for NLL loss.
        Ref: Regularizing Neural Networks by Penalizing Confident Output
        Distributions. https://arxiv.org/abs/1701.06548
    reduction : str
        One of 'mean', 'batch', 'batchmean', 'none' where 'mean' returns a
        single value and 'batch' returns one per item in the batch and
        'batchmean' is sum / batch_size and 'none' returns all.
    """
    mask = torch.ones_like(targets)
    if length is not None:
        length_mask = length_to_mask(
            length * targets.shape[1], max_len=targets.shape[1],
        )

        # Handle any dimensionality of input
        while len(length_mask.shape) < len(mask.shape):
            length_mask = length_mask.unsqueeze(-1)
        length_mask = length_mask.type(mask.dtype)
        mask *= length_mask

    # Compute, then reduce loss
    loss = loss_fn(predictions, targets) * mask
    N = loss.size(0)
    if reduction == "mean":
        loss = loss.sum() / torch.sum(mask)
    elif reduction == "batchmean":
        loss = loss.sum() / N
    elif reduction == "batch":
        loss = loss.reshape(N, -1).sum(1) / mask.reshape(N, -1).sum(1)

    if label_smoothing == 0:
        return loss
    else:
        loss_reg = torch.mean(predictions, dim=1) * mask
        if reduction == "mean":
            loss_reg = torch.sum(loss_reg) / torch.sum(mask)
        elif reduction == "batchmean":
            loss_reg = torch.sum(loss_reg) / targets.shape[0]
        elif reduction == "batch":
            loss_reg = loss_reg.sum(1) / mask.sum(1)

        return -label_smoothing * loss_reg + (1 - label_smoothing) * loss


def get_si_snr_with_pitwrapper(source, estimate_source):
    """This function wraps si_snr calculation with the speechbrain pit-wrapper.

    Arguments:
    ---------
    source: [B, T, C],
        Where B is the batch size, T is the length of the sources, C is
        the number of sources the ordering is made so that this loss is
        compatible with the class PitWrapper.

    estimate_source: [B, T, C]
        The estimated source.

    Example:
    ---------
    >>> x = torch.arange(600).reshape(3, 100, 2)
    >>> xhat = x[:, :, (1, 0)]
    >>> si_snr = -get_si_snr_with_pitwrapper(x, xhat)
    >>> print(si_snr)
    tensor([135.2284, 135.2284, 135.2284])
    """

    pit_si_snr = PitWrapper(cal_si_snr)
    loss, perms = pit_si_snr(source, estimate_source)

    return loss, perms


def get_snr_with_pitwrapper(source, estimate_source):
    """This function wraps si_snr calculation with the speechbrain pit-wrapper.
    Arguments:
    ---------
    source: [B, T, E, C],
        Where B is the batch size, T is the length of the sources, E is binaural channels, C is the number of sources
        the ordering is made so that this loss is compatible with the class PitWrapper.
    estimate_source: [B, T, E, C]
        The estimated source.
    """

    pit_snr = PitWrapper(cal_snr)
    loss, perms = pit_snr(source, estimate_source)

    return loss


def cal_si_snr(source, estimate_source):
    """Calculate SI-SNR.

    Arguments:
    ---------
    source: [T, B, C],
        Where B is batch size, T is the length of the sources, C is the number of sources
        the ordering is made so that this loss is compatible with the class PitWrapper.

    estimate_source: [T, B, C]
        The estimated source.

    Example:
    ---------
    >>> import numpy as np
    >>> x = torch.Tensor([[1, 0], [123, 45], [34, 5], [2312, 421]])
    >>> xhat = x[:, (1, 0)]
    >>> x = x.unsqueeze(-1).repeat(1, 1, 2)
    >>> xhat = xhat.unsqueeze(1).repeat(1, 2, 1)
    >>> si_snr = -cal_si_snr(x, xhat)
    >>> print(si_snr)
    tensor([[[ 25.2142, 144.1789],
             [130.9283,  25.2142]]])
    """
    EPS = 1e-8
    assert source.size() == estimate_source.size()
    device = estimate_source.device.type

    source_lengths = torch.tensor(
        [estimate_source.shape[0]] * estimate_source.shape[-2], device=device
    )
    mask = get_mask(source, source_lengths)
    estimate_source *= mask

    num_samples = (
        source_lengths.contiguous().reshape(1, -1, 1).float()
    )  # [1, B, 1]
    mean_target = torch.sum(source, dim=0, keepdim=True) / num_samples
    mean_estimate = (
        torch.sum(estimate_source, dim=0, keepdim=True) / num_samples
    )
    zero_mean_target = source - mean_target
    zero_mean_estimate = estimate_source - mean_estimate
    # mask padding position along T
    zero_mean_target *= mask
    zero_mean_estimate *= mask

    # Step 2. SI-SNR with PIT
    # reshape to use broadcast
    s_target = zero_mean_target  # [T, B, C]
    s_estimate = zero_mean_estimate  # [T, B, C]
    # s_target = <s', s>s / ||s||^2
    dot = torch.sum(s_estimate * s_target, dim=0, keepdim=True)  # [1, B, C]
    s_target_energy = (
        torch.sum(s_target ** 2, dim=0, keepdim=True) + EPS
    )  # [1, B, C]
    proj = dot * s_target / s_target_energy  # [T, B, C]
    # e_noise = s' - s_target
    e_noise = s_estimate - proj  # [T, B, C]
    # SI-SNR = 10 * log_10(||s_target||^2 / ||e_noise||^2)
    si_snr_beforelog = torch.sum(proj ** 2, dim=0) / (
        torch.sum(e_noise ** 2, dim=0) + EPS
    )
    si_snr = 10 * torch.log10(si_snr_beforelog + EPS)  # [B, C]

    return -si_snr.unsqueeze(0)


def cal_snr(source, estimate_source):
    """Calculate binaural channel SNR.
    Arguments:
    ---------
    source: [T, E, B, C],
        Where B is batch size, T is the length of the sources, E is binaural channels, C is the number of sources
        the ordering is made so that this loss is compatible with the class PitWrapper.
    estimate_source: [T, E, B, C]
        The estimated source.
    """
    EPS = 1e-8
    assert source.size() == estimate_source.size()
    device = estimate_source.device.type

    source_lengths = torch.tensor(
        [estimate_source.shape[0]] * estimate_source.shape[-2], device=device
    )
    mask = get_mask(source, source_lengths)  # [T, E, 1]
    estimate_source *= mask

    num_samples = (
        source_lengths.contiguous().reshape(1, -1, 1).float()
    )  # [1, B, 1]
    mean_target = torch.sum(source, dim=0, keepdim=True) / num_samples
    mean_estimate = (
        torch.sum(estimate_source, dim=0, keepdim=True) / num_samples
    )
    zero_mean_target = source - mean_target
    zero_mean_estimate = estimate_source - mean_estimate
    # mask padding position along T
    zero_mean_target *= mask
    zero_mean_estimate *= mask

    # Step 2. SNR with PIT
    # reshape to use broadcast
    s_target = zero_mean_target  # [T, E, B, C]
    s_estimate = zero_mean_estimate  # [T, E, B, C]
    # SNR = 10 * log_10(||s_target||^2 / ||e_noise||^2)
    # n_dim = [x for x in range(len(s_target.shape)-2)]
    snr_beforelog = torch.sum(s_target ** 2, dim=0) / (
        torch.sum((s_estimate - s_target) ** 2, dim=0) + EPS
    )
    snr = 10 * torch.log10(snr_beforelog + EPS)  # [B, C]

    return -snr.unsqueeze(0)


def get_mask(source, source_lengths):
    """
    Arguments
    ---------
    source : [T, B, C]
    source_lengths : [B]

    Returns
    -------
    mask : [T, B, 1]

    Example:
    ---------
    >>> source = torch.randn(4, 3, 2)
    >>> source_lengths = torch.Tensor([2, 1, 4]).int()
    >>> mask = get_mask(source, source_lengths)
    >>> print(mask)
    tensor([[[1.],
             [1.],
             [1.]],
    <BLANKLINE>
            [[1.],
             [0.],
             [1.]],
    <BLANKLINE>
            [[0.],
             [0.],
             [1.]],
    <BLANKLINE>
            [[0.],
             [0.],
             [1.]]])
    """
    mask = source.new_ones(source.size()[:-1]).unsqueeze(-1).transpose(1, -2)
    B = source.size(-2)
    for i in range(B):
        mask[source_lengths[i] :, i] = 0
    return mask.transpose(-2, 1)


class AngularMargin(nn.Module):
    """
    An implementation of Angular Margin (AM) proposed in the following
    paper: '''Margin Matters: Towards More Discriminative Deep Neural Network
    Embeddings for Speaker Recognition''' (https://arxiv.org/abs/1906.07317)

    Arguments
    ---------
    margin : float
        The margin for cosine similiarity
    scale : float
        The scale for cosine similiarity

    Return
    ---------
    predictions : torch.Tensor

    Example
    -------
    >>> pred = AngularMargin()
    >>> outputs = torch.tensor([ [1., -1.], [-1., 1.], [0.9, 0.1], [0.1, 0.9] ])
    >>> targets = torch.tensor([ [1., 0.], [0., 1.], [ 1., 0.], [0.,  1.] ])
    >>> predictions = pred(outputs, targets)
    >>> predictions[:,0] > predictions[:,1]
    tensor([ True, False,  True, False])
    """

    def __init__(self, margin=0.0, scale=1.0):
        super(AngularMargin, self).__init__()
        self.margin = margin
        self.scale = scale

    def forward(self, outputs, targets):
        """Compute AM between two tensors

        Arguments
        ---------
        outputs : torch.Tensor
            The outputs of shape [N, C], cosine similarity is required.
        targets : torch.Tensor
            The targets of shape [N, C], where the margin is applied for.

        Return
        ---------
        predictions : torch.Tensor
        """
        outputs = outputs - self.margin * targets
        return self.scale * outputs


class AdditiveAngularMargin(AngularMargin):
    """
    An implementation of Additive Angular Margin (AAM) proposed
    in the following paper: '''Margin Matters: Towards More Discriminative Deep
    Neural Network Embeddings for Speaker Recognition'''
    (https://arxiv.org/abs/1906.07317)

    Arguments
    ---------
    margin : float
        The margin for cosine similiarity.
    scale: float
        The scale for cosine similiarity.

    Returns
    -------
    predictions : torch.Tensor
        Tensor.
    Example
    -------
    >>> outputs = torch.tensor([ [1., -1.], [-1., 1.], [0.9, 0.1], [0.1, 0.9] ])
    >>> targets = torch.tensor([ [1., 0.], [0., 1.], [ 1., 0.], [0.,  1.] ])
    >>> pred = AdditiveAngularMargin()
    >>> predictions = pred(outputs, targets)
    >>> predictions[:,0] > predictions[:,1]
    tensor([ True, False,  True, False])
    """

    def __init__(self, margin=0.0, scale=1.0, easy_margin=False):
        super(AdditiveAngularMargin, self).__init__(margin, scale)
        self.easy_margin = easy_margin

        self.cos_m = math.cos(self.margin)
        self.sin_m = math.sin(self.margin)
        self.th = math.cos(math.pi - self.margin)
        self.mm = math.sin(math.pi - self.margin) * self.margin

    def forward(self, outputs, targets):
        """
        Compute AAM between two tensors

        Arguments
        ---------
        outputs : torch.Tensor
            The outputs of shape [N, C], cosine similarity is required.
        targets : torch.Tensor
            The targets of shape [N, C], where the margin is applied for.

        Return
        ---------
        predictions : torch.Tensor
        """
        cosine = outputs.float()
        cosine = torch.clamp(cosine, -1 + 1e-7, 1 - 1e-7)
        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m  # cos(theta + m)
        if self.easy_margin:
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
        outputs = (targets * phi) + ((1.0 - targets) * cosine)
        return self.scale * outputs


class LogSoftmaxWrapper(nn.Module):
    """
    Arguments
    ---------
    Returns
    ---------
    loss : torch.Tensor
        Learning loss
    predictions : torch.Tensor
        Log probabilities
    Example
    -------
    >>> outputs = torch.tensor([ [1., -1.], [-1., 1.], [0.9, 0.1], [0.1, 0.9] ])
    >>> outputs = outputs.unsqueeze(1)
    >>> targets = torch.tensor([ [0], [1], [0], [1] ])
    >>> log_prob = LogSoftmaxWrapper(nn.Identity())
    >>> loss = log_prob(outputs, targets)
    >>> 0 <= loss < 1
    tensor(True)
    >>> log_prob = LogSoftmaxWrapper(AngularMargin(margin=0.2, scale=32))
    >>> loss = log_prob(outputs, targets)
    >>> 0 <= loss < 1
    tensor(True)
    >>> outputs = torch.tensor([ [1., -1.], [-1., 1.], [0.9, 0.1], [0.1, 0.9] ])
    >>> log_prob = LogSoftmaxWrapper(AdditiveAngularMargin(margin=0.3, scale=32))
    >>> loss = log_prob(outputs, targets)
    >>> 0 <= loss < 1
    tensor(True)
    """

    def __init__(self, loss_fn):
        super(LogSoftmaxWrapper, self).__init__()
        self.loss_fn = loss_fn
        self.criterion = torch.nn.KLDivLoss(reduction="sum")

    def forward(self, outputs, targets, length=None):
        """
        Arguments
        ---------
        outputs : torch.Tensor
            Network output tensor, of shape
            [batch, 1, outdim].
        targets : torch.Tensor
            Target tensor, of shape [batch, 1].

        Returns
        -------
        loss: torch.Tensor
            Loss for current examples.
        """
        outputs = outputs.squeeze(1)
        targets = targets.squeeze(1)
        targets = F.one_hot(targets.long(), outputs.shape[1]).float()
        try:
            predictions = self.loss_fn(outputs, targets)
        except TypeError:
            predictions = self.loss_fn(outputs)

        predictions = F.log_softmax(predictions, dim=1)
        loss = self.criterion(predictions, targets) / targets.sum()
        return loss


def ctc_loss_kd(log_probs, targets, input_lens, blank_index, device):
    """Knowledge distillation for CTC loss.

    Reference
    ---------
    Distilling Knowledge from Ensembles of Acoustic Models for Joint CTC-Attention End-to-End Speech Recognition.
    https://arxiv.org/abs/2005.09310

    Arguments
    ---------
    log_probs : torch.Tensor
        Predicted tensor from student model, of shape [batch, time, chars].
    targets : torch.Tensor
        Predicted tensor from single teacher model, of shape [batch, time, chars].
    input_lens : torch.Tensor
        Length of each utterance.
    blank_index : int
        The location of the blank symbol among the character indexes.
    device : str
        Device for computing.
    """
    scores, predictions = torch.max(targets, dim=-1)

    pred_list = []
    pred_len_list = []
    for j in range(predictions.shape[0]):
        # Getting current predictions
        current_pred = predictions[j]

        actual_size = (input_lens[j] * log_probs.shape[1]).round().int()
        current_pred = current_pred[0:actual_size]
        current_pred = filter_ctc_output(
            list(current_pred.cpu().numpy()), blank_id=blank_index
        )
        current_pred_len = len(current_pred)
        pred_list.append(current_pred)
        pred_len_list.append(current_pred_len)

    max_pred_len = max(pred_len_list)
    for j in range(predictions.shape[0]):
        diff = max_pred_len - pred_len_list[j]
        for n in range(diff):
            pred_list[j].append(0)

    # generate soft label of teacher model
    fake_lab = torch.from_numpy(np.array(pred_list))
    fake_lab.to(device)
    fake_lab = fake_lab.int()
    fake_lab_lengths = torch.from_numpy(np.array(pred_len_list)).int()
    fake_lab_lengths.to(device)

    input_lens = (input_lens * log_probs.shape[1]).round().int()
    log_probs = log_probs.transpose(0, 1)
    return torch.nn.functional.ctc_loss(
        log_probs,
        fake_lab,
        input_lens,
        fake_lab_lengths,
        blank_index,
        zero_infinity=True,
    )


def ce_kd(inp, target):
    """Simple version of distillation for cross-entropy loss.

    Arguments
    ---------
    inp : torch.Tensor
        The probabilities from student model, of shape [batch_size * length, feature]
    target : torch.Tensor
        The probabilities from teacher model, of shape [batch_size * length, feature]
    """
    return (-target * inp).sum(1)


def nll_loss_kd(
    probabilities, targets, rel_lab_lengths,
):
    """Knowledge distillation for negative log-likelihood loss.

    Reference
    ---------
    Distilling Knowledge from Ensembles of Acoustic Models for Joint CTC-Attention End-to-End Speech Recognition.
    https://arxiv.org/abs/2005.09310

    Arguments
    ---------
    probabilities : torch.Tensor
        The predicted probabilities from the student model.
        Format is [batch, frames, p]
    targets : torch.Tensor
        The target probabilities from the teacher model.
        Format is [batch, frames, p]
    rel_lab_lengths : torch.Tensor
        Length of each utterance, if the frame-level loss is desired.

    Example
    -------
    >>> probabilities = torch.tensor([[[0.8, 0.2], [0.2, 0.8]]])
    >>> targets = torch.tensor([[[0.9, 0.1], [0.1, 0.9]]])
    >>> rel_lab_lengths = torch.tensor([1.])
    >>> nll_loss_kd(probabilities, targets, rel_lab_lengths)
    tensor(-0.7400)
    """
    # Getting the number of sentences in the minibatch
    N_snt = probabilities.shape[0]

    # Getting the maximum length of label sequence
    max_len = probabilities.shape[1]

    # Getting the label lengths
    lab_lengths = torch.round(rel_lab_lengths * targets.shape[1]).int()

    # Reshape to [batch_size * length, feature]
    prob_curr = probabilities.reshape(N_snt * max_len, probabilities.shape[-1])

    # Generating mask
    mask = length_to_mask(
        lab_lengths, max_len=max_len, dtype=torch.float, device=prob_curr.device
    )

    # Reshape to [batch_size * length, feature]
    lab_curr = targets.reshape(N_snt * max_len, targets.shape[-1])

    loss = ce_kd(prob_curr, lab_curr)
    # Loss averaging
    loss = torch.sum(loss.reshape(N_snt, max_len) * mask) / torch.sum(mask)
    return loss


class ContrastiveLoss(nn.Module):
    """Contrastive loss as used in wav2vec2.

    Reference
    ---------
    wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations
    https://arxiv.org/abs/2006.11477

    Arguments
    ---------
    logit_temp : torch.Float
        A temperature to devide the logits.

    """

    def __init__(self, logit_temp):
        super().__init__()
        self.logit_temp = logit_temp

    def forward(self, x, y, negs):
        """
        Arguments
        ----------
        x : torch.Tensor
            Encoded embeddings with shape (B, T, C).
        y : torch.Tensor
            Feature extractor target embeddings with shape (B, T, C).
        negs : torch.Tensor
            Negative embeddings from feature extractor with shape (N, B, T, C)
        where N is number of negatives. Can be obtained with our sample_negatives
        function (check in lobes/wav2vec2).
        """
        neg_is_pos = (y == negs).all(-1)
        y = y.unsqueeze(0)
        target_and_negatives = torch.cat([y, negs], dim=0)
        logits = torch.cosine_similarity(
            x.float(), target_and_negatives.float(), dim=-1
        ).type_as(x)

        if neg_is_pos.any():
            logits[1:][neg_is_pos] = float("-inf")
        # N, B, T -> T, B, N -> T*B, N
        logits = logits.transpose(0, 2).reshape(-1, logits.size(0))

        targets = torch.zeros(
            (logits.size(0)), dtype=torch.long, device=logits.device
        )
        loss = F.cross_entropy(
            logits / self.logit_temp, targets, reduction="sum"
        )
        accuracy = torch.sum(logits.argmax(-1) == 0) / (
            logits.numel() / logits.size(-1)
        )
        return loss, accuracy

speechbrain/nnet/losses.py
TzuchengChang/NASS
from packaging import version
import torch
from torch import nn


class PatchNCELoss(nn.Module):
    def __init__(self, opt):
        super().__init__()
        self.opt = opt
        self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none')
        self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool

    def forward(self, feat_q, feat_p, feat_n):
        num_patches = feat_q.shape[0]
        dim = feat_q.shape[1]
        feat_p = feat_p.detach()
        feat_n = feat_n.detach()

        l_pos = torch.bmm(
            feat_q.view(num_patches, 1, -1), feat_p.view(num_patches, -1, 1))
        l_pos = l_pos.view(num_patches, 1)

        if self.opt['nce_includes_all_negatives_from_minibatch']:
            batch_dim_for_bmm = 1
        else:
            batch_dim_for_bmm = self.opt['batch_size']

        feat_q = feat_q.view(batch_dim_for_bmm, -1, dim)
        feat_n = feat_n.view(batch_dim_for_bmm, -1, dim)
        npatches = feat_q.size(1)
        l_neg_curbatch = torch.bmm(feat_q, feat_n.transpose(2, 1))

        diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :]
        l_neg_curbatch.masked_fill_(diagonal, -10.0)
        l_neg = l_neg_curbatch.view(-1, npatches)

        out = torch.cat((l_pos, l_neg), dim=1) / self.opt['nce_T']

        loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long,
                                                        device=feat_q.device))

        return loss

speechbrain/nnet/patchnce.py
TzuchengChang/NASS
../speaker_id/custom_model.py
templates/hyperparameter_optimization_speaker_id/custom_model.py
TzuchengChang/NASS
../speaker_id/mini_librispeech_prepare.py
templates/hyperparameter_optimization_speaker_id/mini_librispeech_prepare.py
TzuchengChang/NASS
""" Availing scripts for testing to be imported, i.e., code in tests/utils.
"""

tests/__init__.py
TzuchengChang/NASS
"""Library for the HuggingFace (HF) repositories.

Authors
 * Mirco Ravanelli 2022
 * Andreas Nautsch 2022
"""
import os
import csv
from speechbrain.utils.data_utils import download_file


def run_HF_check(
    recipe_folder="tests/recipes", field="HF_repo", output_folder="HF_repos",
):
    """Checks if the code reported in the readme files of the HF repository is
    runnable. Note: the tests run the code marked as python in the readme file.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder with csv recipe files summarizing all the recipes in the repo.
    field: string
        Field of the csv recipe file containing the links to HF repos.
    output_folder: path
        Where to download the HF readme files.

    Returns
    ---------
    check: True
        True if all the code runs, False otherwise.
    """
    check = True
    for recipe_csvfile in os.listdir(recipe_folder):
        # Detect list of HF repositories
        HF_repos = repo_list(os.path.join(recipe_folder, recipe_csvfile), field)

        # Set up output folder
        os.makedirs(output_folder, exist_ok=True)
        os.chdir(output_folder)

        # Checking all detected repos
        for repo in HF_repos:
            if not (check_repo(repo)):
                check = False
    return check


def repo_list(recipe_folder="tests/recipes", field="HF_repo"):
    """Get the list of HF recipes in the csv recipe file.

    Arguments
    ---------
    recipe_folder: path
        Path of the fodler with csv recipe files summarizing all the recipes in the repo.
    field: string
        Field of the csv recipe file containing the links to HF repos.

    Returns
    ---------
    HF_repos: list
        List of the detected HF repos.
    """
    HF_repos = []

    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvf:
            reader = csv.DictReader(csvf, delimiter=",", skipinitialspace=True)
            for row in reader:
                if len(row[field]) > 0:
                    repos = row[field].split(" ")
                    for repo in repos:
                        HF_repos.append(repo)
    HF_repos = set(HF_repos)
    return HF_repos


def check_repo(HF_repo):
    """Runs the code reported in the README file of the given HF_repo. It checks
    if the code runs without errors.

    Arguments
    ---------
    HF_repo: string
        URL of the HF repository to check.

    Returns
    ---------
    check: bool
        True if all the code runs, False otherwise.
    """
    exp_name = os.path.basename(HF_repo)
    readme_file = HF_repo + "/raw/main/README.md"

    dest_file = exp_name + ".md"
    download_file(readme_file, dest_file)

    code_snippets = []
    code = []
    flag = False
    with open(dest_file, "r") as f:
        for line in f:
            if "```python" in line:
                flag = True
                code = []
            elif "```\n" in line and flag:
                flag = False
                code_snippets.append(code)
            elif flag:
                if len(line.strip()) > 0:
                    code.append(line)
                    print(line)

    for code in code_snippets:
        try:
            exec("\n".join(code))
        except Exception as e:
            print("\t" + str(e))
            check = False
            print("\tERROR: cannot run code snippet in %s" % (HF_repo))
    return check

tests/consistency/test_HF_repo.py
TzuchengChang/NASS
"""Tests for checking the docstrings of functions and classes.

Authors
 * Mirco Ravanelli 2022
"""
from tests.utils.check_docstrings import check_docstrings


def test_recipe_list(base_folder="."):
    check_folders = ["speechbrain", "tools", "templates"]
    assert check_docstrings(base_folder, check_folders)

tests/consistency/test_docstrings.py
TzuchengChang/NASS
"""Tests for checking the recipes and their files.

Authors
 * Mirco Ravanelli 2022
"""
import os
import csv
from speechbrain.utils.data_utils import get_all_files, get_list_from_csv


__skip_list = ["README.md", "setup"]


def test_recipe_list(
    search_folder="recipes",
    hparam_ext=[".yaml"],
    hparam_field="Hparam_file",
    recipe_folder="tests/recipes",
    flags_field="test_debug_flags",
    avoid_list=[
        "/models/",
        "/results/",
        "recipes/Voicebank/MTL/CoopNet/hparams/logger.yaml",
        "recipes/LibriParty/generate_dataset/dataset.yaml",
        "hpopt.yaml",
    ],
):
    """This test checks if all the all hparam file of all the recipes are listed
    in the csv recipe file.

    Arguments
    ---------
    search_folder: path
        The path where to search the hparam files.
    hparam_ext: list
        The list containing the extensions of hparam files.
    hparam_field: str
        Field of the csv file where the path of the hparam file is reported.
    recipe_folder: path
        Path of the folder containing csv recipe files.
    flags_field: str
        Field of the csv file where the debug flags are stated (for data flow testing).
    avoid_list: list
        List of files for which this check must be avoided.

    Returns
    ---------
    bool:
        True if the test passes, False otherwise.

    """
    all_diffs_zero = True
    all_with_flags = True
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue
        dataset = os.path.splitext(recipe_csvfile)[0]
        hparam_lst = get_all_files(
            os.path.join(search_folder, dataset),
            match_and=hparam_ext,
            exclude_or=avoid_list,
        )
        recipe_lst = get_list_from_csv(
            os.path.join(recipe_folder, recipe_csvfile), field=hparam_field
        )
        diff_lst = list(set(hparam_lst) - set(recipe_lst))

        for file in diff_lst:
            print(
                "\tERROR: The file %s is not listed in %s. Please add it. \
                    For more info see tests/consistency/README.md"
                % (file, recipe_csvfile)
            )

        all_diffs_zero &= len(diff_lst) == 0

        flags_lst = get_list_from_csv(
            os.path.join(recipe_folder, recipe_csvfile), flags_field
        )
        for flags in flags_lst:
            if not flags:
                all_with_flags = False
                print(f"\tERROR: {flags_field} are missing in {recipe_csvfile}")

    assert all_diffs_zero
    assert all_with_flags


def test_recipe_files(
    recipe_folder="tests/recipes",
    fields=["Script_file", "Hparam_file", "Data_prep_file", "Readme_file"],
):
    """This test checks if the files listed in the recipe csv file exist.

    Arguments
    ---------.
    recipe_folder: path
        Path of the folder containing csv recipe files.
    fields: list
        Fields of the csv recipe file to check.

    Returns
    ---------
    check: bool
        True if the test passes, False otherwise.
    """
    check = True
    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue
        for field in fields:
            lst = get_list_from_csv(
                os.path.join(recipe_folder, recipe_csvfile), field=field
            )
            lst = filter(None, lst)
            for files in lst:
                files = files.split(" ")
                files = filter(None, files)
                for file in files:
                    if not (os.path.exists(file.strip())):
                        print(
                            "\tERROR: The file %s listed in %s does not exist!"
                            % (file, recipe_csvfile)
                        )
                        check = False
    assert check


def test_mandatory_files(
    recipe_folder="tests/recipes",
    must_exist=["Script_file", "Hparam_file", "Readme_file"],
):
    """This test checks if all the recipes contain the specified mandatory files.

    Arguments
    ---------.
    recipe_folder: path
        Path of the folder containing csv recipe files.
    must_exist: list
        List of the fields of the csv recipe file that must contain valid paths.

    Returns
    ---------
    check: bool
        True if the test passes, False otherwise.
    """

    check = True

    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue
        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvf:
            reader = csv.DictReader(csvf, delimiter=",", skipinitialspace=True)
            for row_id, row in enumerate(reader):
                for field in must_exist:
                    if not (os.path.exists(row[field].strip())):
                        print(
                            "\tERROR: The recipe %s does not contain a %s. Please add it to %s!"
                            % (row_id, field, recipe_csvfile)
                        )
                        check = False
    assert check


def test_README_links(
    recipe_folder="tests/recipes",
    readme_field="Readme_file",
    must_link=["Result_url", "HF_repo"],
):
    """This test checks if the README file contains the correct GDRIVE and HF repositories.

    Arguments
    ---------.
    recipe_folder: path
        Path of the folder containing csv recipe files.
    readme_field: string
        Field of the csv recipe file that contains the path to the readme file.
    must_link : list
        Fields that contains the paths that must be linked in the readme file.

    Returns
    ---------
    check: bool
        True if the test passes, False otherwise.
    """
    check = True

    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue
        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvf:
            reader = csv.DictReader(csvf, delimiter=",", skipinitialspace=True)
            for row in reader:
                with open(row[readme_field].strip()) as readmefile:
                    content = readmefile.read()
                    for field in must_link:
                        links = row[field].strip().split(" ")
                        for link in links:
                            if len(link) == 0:
                                continue
                            if not (link in content):
                                print(
                                    "\tERROR: The link to %s does not exist in %s. Please add it to %s!"
                                    % (link, row[readme_field], recipe_csvfile)
                                )
                                check = False
    assert check

tests/consistency/test_recipe.py
TzuchengChang/NASS
"""Consistency check between yaml files and script files.

Authors
 * Mirco Ravanelli 2022
"""
import os
import csv
from tests.consistency.test_recipe import __skip_list
from tests.utils.check_yaml import check_yaml_vs_script


def test_yaml_script_consistency(recipe_folder="tests/recipes"):
    """This test checks the consistency between yaml files (used to specify
    hyperparameters) and script files (that implement the training recipe).

    Arguments
    ---------
    recipe_folder : path
        Path of the folder with csv files containing the training scripts with their coupled
        yaml files (with colums called 'Hparam_file', 'Script_file', 'Data_prep_file')
    """

    # Use this list to itemize special yaml for which we do not have to test
    avoid_check = []

    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue
        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvfile:
            check = True
            reader = csv.DictReader(
                csvfile, delimiter=",", skipinitialspace=True
            )
            for row in reader:

                # Avoid checks
                if row["Hparam_file"] in avoid_check:
                    continue

                # Check yaml-script consistency
                if not (
                    check_yaml_vs_script(row["Hparam_file"], row["Script_file"])
                ):
                    check = False

                # Check module variables
                # if not (check_module_vars(row["Hparam_file"], row["Script_file"])):
                #    check = False

    assert check

tests/consistency/test_yaml.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a CTC-based speech recognizer on a tiny dataset.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phonemes.  A greedy search is used on
top of the output probabilities.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class CTCBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.modules.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x = self.modules.model(feats)
        x = self.modules.lin(x)
        outputs = self.hparams.softmax(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the CTC loss."
        predictions, lens = predictions
        phns, phn_lens = batch.phn_encoded
        loss = self.hparams.compute_cost(predictions, phns, lens, phn_lens)

        if stage != sb.Stage.TRAIN:
            seq = sb.decoders.ctc_greedy_decode(
                predictions, lens, blank_id=self.hparams.blank_index
            )
            self.per_metrics.append(batch.id, seq, phns, target_len=phn_lens)

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss

        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)

        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides("phn_list", "phn_encoded")
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.insert_blank(index=hparams["blank_index"])
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "phn_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR/"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ctc_brain = CTCBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ctc_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ctc_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ctc_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_CTC/example_asr_ctc_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a CTC-based speech recognizer on a tiny dataset.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phonemes.  A greedy search is used on
top of the output probabilities.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class CTCBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.modules.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x = self.modules.model(feats)
        x = self.modules.lin(x)
        outputs = self.hparams.softmax(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the CTC loss."
        predictions, lens = predictions
        phns, phn_lens = batch.phn_encoded
        loss = self.hparams.compute_cost(predictions, phns, lens, phn_lens)

        if stage != sb.Stage.TRAIN:
            seq = sb.decoders.ctc_greedy_decode(
                predictions, lens, blank_id=self.hparams.blank_index
            )
            self.per_metrics.append(batch.id, seq, phns, target_len=phn_lens)

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss

        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)

        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides("phn_list", "phn_encoded")
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.insert_blank(index=hparams["blank_index"])
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "phn_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams_complex_net.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ctc_brain = CTCBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ctc_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ctc_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ctc_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_CTC/example_asr_ctc_experiment_complex_net.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a CTC-based speech recognizer on a tiny dataset.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phonemes.  A greedy search is used on
top of the output probabilities.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class CTCBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.modules.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x = self.modules.model(feats)
        x = self.modules.lin(x)
        outputs = self.hparams.softmax(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the CTC loss."
        predictions, lens = predictions
        phns, phn_lens = batch.phn_encoded
        loss = self.hparams.compute_cost(predictions, phns, lens, phn_lens)

        if stage != sb.Stage.TRAIN:
            seq = sb.decoders.ctc_greedy_decode(
                predictions, lens, blank_id=self.hparams.blank_index
            )
            self.per_metrics.append(batch.id, seq, phns, target_len=phn_lens)

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss

        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)

        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides("phn_list", "phn_encoded")
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.insert_blank(index=hparams["blank_index"])
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "phn_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams_quaternion_net.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ctc_brain = CTCBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ctc_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ctc_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ctc_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_CTC/example_asr_ctc_experiment_quaternion_net.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a RNNT-based speech recognizer on a tiny dataset.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phonemes.  A beamsearch is used on
top of the output probabilities.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""
import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class TransducerBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.modules.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)

        # Transcription network: input-output dependency
        TN_output = self.modules.enc(feats)
        TN_output = self.modules.enc_lin(TN_output)

        # Prediction network: output-output dependency
        targets, target_lens = batch.phn_encoded_bos
        PN_output = self.modules.emb(targets)
        PN_output, _ = self.modules.dec(PN_output)
        PN_output = self.modules.dec_lin(PN_output)

        # Joint the networks
        joint = self.modules.Tjoint(
            TN_output.unsqueeze(2), PN_output.unsqueeze(1),
        )
        outputs = self.modules.output(joint)
        outputs = self.hparams.log_softmax(outputs)
        if stage == sb.Stage.TRAIN:
            return outputs, lens
        else:
            hyps, scores, _, _ = self.hparams.searcher(TN_output)
            return outputs, lens, hyps

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the CTC loss."
        phns, phn_lens = batch.phn_encoded

        if stage == sb.Stage.TRAIN:
            predictions, lens = predictions
        else:
            predictions, lens, seq = predictions
            self.per_metrics.append(batch.id, seq, phns, target_len=phn_lens)

        loss = self.hparams.compute_cost(
            predictions,
            phns.to(self.device).long(),
            lens,
            phn_lens.to(self.device),
        )
        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides(
        "phn_list", "phn_encoded", "phn_encoded_bos"
    )
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded
        phn_encoded_bos = label_encoder.prepend_bos_index(phn_encoded).long()
        yield phn_encoded_bos

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.insert_blank(index=hparams["blank_index"])
    label_encoder.insert_bos_eos(
        bos_index=hparams["bos_index"], eos_label="<bos>"
    )
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets, ["id", "sig", "phn_encoded", "phn_encoded_bos"]
    )
    return train_data, valid_data, label_encoder


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data, label_encoder = data_prep(data_folder, hparams)

    # Trainer initialization
    transducer_brain = TransducerBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    transducer_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    transducer_brain.evaluate(valid_data)

    # Check that model overfits for integration test
    assert transducer_brain.train_loss < 90.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_Transducer/example_asr_transducer_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains an HMM-based aligner with the forward algorithm.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phoneme states.
Given the tiny dataset, the expected behavior is to overfit the training data
(with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class AlignBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x = self.modules.model(feats)
        x = self.modules.lin(x)
        outputs = self.hparams.softmax(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the forward loss."
        predictions, lens = predictions
        phns, phn_lens = batch.phn_encoded
        sum_alpha_T = self.hparams.aligner(
            predictions, lens, phns, phn_lens, "forward"
        )
        loss = -sum_alpha_T.sum()

        if stage != sb.Stage.TRAIN:
            viterbi_scores, alignments = self.hparams.aligner(
                predictions, lens, phns, phn_lens, "viterbi"
            )

        return loss

    def on_stage_end(self, stage, stage_loss, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)

        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides("phn_list", "phn_encoded")
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "phn_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR/"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ali_brain = AlignBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ali_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ali_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ali_brain.train_loss < 350


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_alignment_forward/example_asr_alignment_forward_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains an HMM-based aligner with the Viterbi algorithm.
The encoder is based on a combination of convolutional, recurrent, and
feed-forward networks (CRDNN) that predict phoneme states.
Given the tiny dataset, the expected behavior is to overfit the training data
(with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class AlignBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.modules.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x = self.modules.model(feats)
        x = self.modules.lin(x)
        outputs = self.hparams.softmax(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the forward loss."
        predictions, lens = predictions
        phns, phn_lens = batch.phn_encoded

        prev_alignments = self.hparams.aligner.get_prev_alignments(
            batch.id, predictions, lens, phns, phn_lens
        )
        loss = self.hparams.compute_cost(predictions, prev_alignments)

        if stage != sb.Stage.TRAIN:
            viterbi_scores, alignments = self.hparams.aligner(
                predictions, lens, phns, phn_lens, "viterbi"
            )
            self.hparams.aligner.store_alignments(batch.id, alignments)

        return loss

    def on_stage_end(self, stage, stage_loss, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
            print("Valid loss: %.2f" % stage_loss)
            print("Recalculating and recording alignments...")
            self.evaluate(self.hparams.train_data)


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )

    # The evaluate method of the brain class, needs to align over training data
    hparams["train_data"] = train_data

    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CTCTextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides("phn_list", "phn_encoded")
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        yield phn_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "phn_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR/"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ali_brain = AlignBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ali_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ali_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ali_brain.train_loss < 2.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_alignment_viterbi/example_asr_alignment_viterbi_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a seq2seq attention-based model for speech
recognition on a tiny dataset.  The encoder is based on a combination of
convolutional, recurrent, and feed-forward networks (CRDNN). The decoder is
based on a GRU. A greedy search  is used on top of the output probabilities.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""
import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class seq2seqBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the output probabilities."
        batch = batch.to(self.device)
        wavs, wav_lens = batch.sig
        phns_bos, _ = batch.phn_encoded_bos
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, wav_lens)
        x = self.modules.enc(feats)

        # Prepend bos token at the beginning
        e_in = self.modules.emb(phns_bos)
        h, w = self.modules.dec(e_in, x, wav_lens)
        logits = self.modules.lin(h)
        outputs = self.hparams.softmax(logits)

        if stage != sb.Stage.TRAIN:
            seq, _ = self.hparams.searcher(x, wav_lens)
            return outputs, seq

        return outputs

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the NLL loss."
        if stage == sb.Stage.TRAIN:
            outputs = predictions
        else:
            outputs, seq = predictions

        ids = batch.id
        phns, phn_lens = batch.phn_encoded_eos

        loss = self.hparams.compute_cost(outputs, phns, length=phn_lens)

        if stage != sb.Stage.TRAIN:
            self.per_metrics.append(ids, seq, phns, target_len=phn_lens)

        return loss

    def fit_batch(self, batch):
        """Fits train batches"""
        preds = self.compute_forward(batch, sb.Stage.TRAIN)
        loss = self.compute_objectives(preds, batch, sb.Stage.TRAIN)
        loss.backward()
        if self.check_gradients(loss):
            self.optimizer.step()
        self.optimizer.zero_grad()
        return loss.detach()

    def evaluate_batch(self, batch, stage=sb.Stage.TEST):
        """Evaluates test batches"""
        out = self.compute_forward(batch, stage)
        loss = self.compute_objectives(out, batch, stage)
        return loss.detach()

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        "Gets called when a stage (either training, validation, test) ends."
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.TextEncoder()
    label_encoder.expect_len(hparams["num_labels"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides(
        "phn_list", "phn_encoded_bos", "phn_encoded_eos"
    )
    def text_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = label_encoder.encode_sequence_torch(phn_list)
        phn_encoded_bos = label_encoder.prepend_bos_index(phn_encoded).long()
        yield phn_encoded_bos
        phn_encoded_eos = label_encoder.append_eos_index(phn_encoded).long()
        yield phn_encoded_eos

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.insert_bos_eos(bos_index=hparams["bos_index"])
    label_encoder.update_from_didataset(train_data, output_key="phn_list")
    label_encoder.update_from_didataset(valid_data, output_key="phn_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets, ["id", "sig", "phn_encoded_eos", "phn_encoded_bos"]
    )
    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    seq2seq_brain = seq2seqBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    seq2seq_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    seq2seq_brain.evaluate(valid_data)

    # Check that model overfits for integration test
    assert seq2seq_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/ASR_seq2seq/example_asr_seq2seq_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a grapheme-to-phoneme (G2P) converter
that turns a sequence of characters into a sequence of phonemes. The system uses
a standard attention-based encoder-decoder pipeline. The encoder is based on an
 LSTM, while the decoder is based on a GRU. Greedy search applied on the top of
the output probabilities to detect the final sequence of phonemes. Given the
tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""
import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class seq2seqBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given input chars it computes the phoneme's probabilities"
        batch = batch.to(self.device)
        chars, char_lens = batch.char_encoded
        phns, phn_lens = batch.phn_encoded_bos

        emb_char = self.modules.encoder_emb(chars)
        x, _ = self.modules.enc(emb_char)
        e_in = self.modules.emb(phns)

        h, w = self.modules.dec(e_in, x, char_lens)
        logits = self.modules.lin(h)
        outputs = self.hparams.softmax(logits)

        if stage != sb.Stage.TRAIN:
            seq, _ = self.hparams.searcher(x, char_lens)
            return outputs, seq

        return outputs

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the NLL loss."
        if stage == sb.Stage.TRAIN:
            outputs = predictions
        else:
            outputs, seq = predictions

        phns, phn_lens = batch.phn_encoded_eos
        loss = self.hparams.compute_cost(outputs, phns, length=phn_lens)

        if stage != sb.Stage.TRAIN:
            self.per_metrics.append(batch.id, seq, phns, target_len=phn_lens)

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.per_metrics = self.hparams.per_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        "Gets called when a stage (either training, validation, test) ends."
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID and epoch is not None:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(stage, "PER: %.2f" % self.per_metrics.summarize("error_rate"))


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    char_encoder = sb.dataio.encoder.TextEncoder()
    phn_encoder = sb.dataio.encoder.TextEncoder()

    # 2. Define char pipeline:
    @sb.utils.data_pipeline.takes("char")
    @sb.utils.data_pipeline.provides("char_list", "char_encoded")
    def char_pipeline(char):
        char_list = char.strip().split()
        yield char_list
        char_encoded = char_encoder.encode_sequence_torch(char_list)
        yield char_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, char_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    char_encoder.insert_bos_eos(bos_index=hparams["bos_index"])
    char_encoder.update_from_didataset(train_data, output_key="char_list")
    char_encoder.update_from_didataset(valid_data, output_key="char_list")
    char_encoder.expect_len(hparams["num_chars"])

    # 4. Define char pipeline:
    @sb.utils.data_pipeline.takes("phn")
    @sb.utils.data_pipeline.provides(
        "phn_list", "phn_encoded_bos", "phn_encoded_eos"
    )
    def phn_pipeline(phn):
        phn_list = phn.strip().split()
        yield phn_list
        phn_encoded = phn_encoder.encode_sequence_torch(phn_list)
        phn_encoded_bos = phn_encoder.prepend_bos_index(phn_encoded).long()
        yield phn_encoded_bos
        phn_encoded_eos = phn_encoder.append_eos_index(phn_encoded).long()
        yield phn_encoded_eos

    sb.dataio.dataset.add_dynamic_item(datasets, phn_pipeline)

    # 5. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    phn_encoder.insert_bos_eos(bos_index=hparams["bos_index"])
    phn_encoder.update_from_didataset(train_data, output_key="phn_list")
    phn_encoder.update_from_didataset(valid_data, output_key="phn_list")
    phn_encoder.expect_len(hparams["num_phns"])

    # 6. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets, ["id", "char_encoded", "phn_encoded_eos", "phn_encoded_bos"]
    )
    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    seq2seq_brain = seq2seqBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    seq2seq_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    seq2seq_brain.evaluate(valid_data)

    # Check that model overfits for integration test
    assert seq2seq_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/G2P/example_g2p.py
TzuchengChang/NASS
#!/usr/bin/env/python3

"""This minimal example trains a character-level language model that predicts
the next characters given the previous ones. The system uses a standard
attention-based encoder-decoder pipeline. The encoder is based on a simple LSTM.
Given the tiny dataset, the expected behavior is to overfit the training dataset
(with a validation performance that stays high).
"""
import math
import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class LMBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input chars it computes the next-char probability."
        batch = batch.to(self.device)
        chars, char_lens = batch.char_encoded_bos
        logits = self.modules.model(chars)
        pout = self.hparams.log_softmax(logits)
        return pout

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the NLL loss."
        chars, char_lens = batch.char_encoded_eos
        loss = self.hparams.compute_cost(predictions, chars, length=char_lens)
        return loss

    def on_stage_end(self, stage, stage_loss, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            perplexity = math.e ** stage_loss
            print(stage, "perplexity: %.2f" % perplexity)


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    char_encoder = sb.dataio.encoder.TextEncoder()
    char_encoder.expect_len(hparams["num_chars"])

    # 2. Define char pipeline:
    @sb.utils.data_pipeline.takes("char")
    @sb.utils.data_pipeline.provides(
        "char_list", "char_encoded_bos", "char_encoded_eos"
    )
    def char_pipeline(char):
        char_list = char.strip().split()
        yield char_list
        char_encoded = char_encoder.encode_sequence_torch(char_list)
        char_encoded_bos = char_encoder.prepend_bos_index(char_encoded).long()
        yield char_encoded_bos
        char_encoded_eos = char_encoder.append_eos_index(char_encoded).long()
        yield char_encoded_eos

    sb.dataio.dataset.add_dynamic_item(datasets, char_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    char_encoder.insert_bos_eos(bos_index=hparams["bos_index"])
    char_encoder.update_from_didataset(train_data, output_key="char_list")
    char_encoder.update_from_didataset(valid_data, output_key="char_list")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets, ["id", "char_encoded_bos", "char_encoded_eos"]
    )
    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    lm_brain = LMBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    lm_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    lm_brain.evaluate(valid_data)

    # Check that model overfits for integration test
    assert lm_brain.train_loss < 0.15


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/LM_RNN/example_lm_rnn_experiment.py
TzuchengChang/NASS
#!/usr/bin/python
import os
import pickle
import numpy
from numpy import linalg as LA
from speechbrain.processing.PLDA_LDA import StatObject_SB  # noqa F401
from speechbrain.processing.PLDA_LDA import PLDA
from speechbrain.processing.PLDA_LDA import Ndx
from speechbrain.processing.PLDA_LDA import fast_PLDA_scoring


# Load params file
experiment_dir = os.path.dirname(os.path.abspath(__file__))
data_folder = "/../../samples/PLDA/"
data_folder = os.path.abspath(experiment_dir + data_folder)

# Xvectors stored as StatObject_SB
train_file = data_folder + "/train_stat_xvect.pkl"
enrol_file = data_folder + "/enrol_stat_xvect.pkl"
test_file = data_folder + "/test_stat_xvect.pkl"
scores_file = data_folder + "/expected_plda_scores.pkl"

# Load Train
with open(train_file, "rb") as input:
    train_obj = pickle.load(input)

# Load Enrol
with open(enrol_file, "rb") as input:
    enrol_obj = pickle.load(input)

# Load Test
with open(test_file, "rb") as input:
    test_obj = pickle.load(input)

print("Training PLDA...")
plda = PLDA()
plda.plda(train_obj)

# Preparing Ndx map
models = enrol_obj.modelset
testsegs = test_obj.modelset
ndx_obj = Ndx(models=models, testsegs=testsegs)

# PLDA scoring between enrol and test
scores_plda = fast_PLDA_scoring(
    enrol_obj, test_obj, ndx_obj, plda.mean, plda.F, plda.Sigma
)
print("PLDA score matrix: (Rows: Enrol, Columns: Test)")
print(scores_plda.scoremat)

with open(scores_file, "rb") as input:
    expected_score_matrix = pickle.load(input)

print("Expected scores:\n", expected_score_matrix)

# Ensuring the scores are proper (for integration test)
dif = numpy.subtract(expected_score_matrix, scores_plda.scoremat)
f_norm = LA.norm(dif, ord="fro")


# Integration test: Ensure we get same score matrix
def test_error():
    assert f_norm < 0.1

tests/integration/PLDA/example_plda_experiment.py
TzuchengChang/NASS
"""This minimal example trains a Voice Activity Detector (VAD) on a tiny dataset.
The network is based on a LSTM with a linear transformation on the top of that.
The system is trained with the binary cross-entropy metric.
"""

import os
import torch
import numpy as np
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class VADBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the binary probability."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x, _ = self.modules.rnn(feats)
        outputs = self.modules.lin(x)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage=True):
        "Given the network predictions and targets computed the binary CE"
        predictions, lens = predictions

        targets, lens = batch.target
        targets = targets.to(predictions.device)
        predictions = predictions[:, : targets.shape[-1], 0]
        loss = self.hparams.compute_BCE_cost(
            torch.nn.BCEWithLogitsLoss(reduction="none"),
            predictions,
            targets,
            lens,
        )

        # compute metrics
        self.binary_metrics.append(
            batch.id, torch.sigmoid(predictions), targets
        )

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        self.binary_metrics = sb.utils.metric_stats.BinaryMetricStats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
            train_summary = self.binary_metrics.summarize(threshold=0.5)

            print("Epoch %d completed" % epoch)
            print("Train loss: %.4f" % stage_loss)
            print("Train Precision: %.2f" % train_summary["precision"])
            print("Train Recall: %.2f" % train_summary["recall"])


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=os.path.join(data_folder, "../annotation/VAD_train.json"),
        replacements={"data_folder": data_folder},
    )

    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=os.path.join(data_folder, "../annotation/VAD_dev.json"),
        replacements={"data_folder": data_folder},
    )

    datasets = [train_data, valid_data]

    # 1. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    # 2. vad targets creation from annotated speech boundaries
    @sb.utils.data_pipeline.takes("speech")
    @sb.utils.data_pipeline.provides("target")
    def vad_targets(string, hparams=hparams):
        if len(string) > 0:
            boundaries = string.split(" ")
            # we group by two
            # 0.01 is 10 ms hop size ...
            boundaries = [int(float(x) / 0.01) for x in boundaries]
            boundaries = list(zip(boundaries[::2], boundaries[1::2]))
        else:
            boundaries = []

        gt = torch.zeros(int(np.ceil(hparams["example_length"] * (1 / 0.01))))

        for indxs in boundaries:
            start, stop = indxs
            gt[start:stop] = 1

        return gt

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)
    sb.dataio.dataset.add_dynamic_item(datasets, vad_targets)

    # 3. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "target"])

    return train_data, valid_data


def main(device="cpu"):

    experiment_dir = os.path.dirname(os.path.abspath(__file__))
    hparams_file = os.path.join(experiment_dir, "hyperparams.yaml")
    data_folder = "/../../samples/VAD"
    data_folder = os.path.abspath(experiment_dir + data_folder)
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Data IO creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    ctc_brain = VADBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    ctc_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    ctc_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert ctc_brain.train_loss < 1.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/VAD/example_vad.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains an autoencoder over speech features. The encoder
is a MLP that transforms the input into a lower-dimensional latent representation.
The decoder is another MLP that predicts the input features. The system is trained
with MSE. Given the tiny dataset, the expected behavior is to overfit the
training data  (with a validation performance that stays high).
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class AutoBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Appling encoder and decoder to the input features"
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)

        encoded = self.modules.linear1(feats)
        encoded = self.hparams.activation(encoded)
        decoded = self.modules.linear2(encoded)

        return decoded

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the MSE loss."
        wavs, lens = batch.sig
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        self.mse_metric.append(batch.id, predictions, feats, lens)
        return self.hparams.compute_cost(predictions, feats, lens)

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        self.mse_metric = self.hparams.loss_tracker()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if self.hparams.use_tensorboard:
            if stage == sb.Stage.TRAIN:
                self.hparams.train_logger.log_stats(
                    {"Epoch": epoch},
                    train_stats={"loss": self.mse_metric.scores},
                )
            elif stage == sb.Stage.VALID:
                self.hparams.train_logger.log_stats(
                    {"Epoch": epoch},
                    valid_stats={"loss": self.mse_metric.scores},
                )
            if stage == sb.Stage.TEST:
                self.hparams.train_logger.log_stats(
                    {}, test_stats={"loss": self.mse_metric.scores}
                )

        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID:
            print("Completed epoch %d" % epoch)
            print("Train loss: %.3f" % self.train_loss)
            print("Valid loss: %.3f" % stage_loss)
        if stage == sb.Stage.TEST:
            print("Test loss: %.3f" % stage_loss)


def data_prep(data_folder):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder)

    # Trainer initialization
    auto_brain = AutoBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    auto_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    auto_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert auto_brain.train_loss < 0.08


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/autoencoder/example_auto_experiment.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a GAN speech enhancement system on a tiny dataset.
The generator and the discriminator are based on convolutional networks.
"""

import torch
import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


class EnhanceGanBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the enhanced signal"
        batch = batch.to(self.device)
        wavs, lens = batch.sig

        noisy = self.hparams.add_noise(wavs, lens).unsqueeze(-1)
        enhanced = self.modules.generator(noisy)

        return enhanced

    def compute_objectives(self, predictions, batch, stage, optim_name=""):
        "Given the network predictions and targets computed the total loss"
        clean_wavs, lens = batch.sig
        batch_size = clean_wavs.size(0)

        # Average the predictions of each time step
        clean_wavs = clean_wavs.unsqueeze(-1)
        real_result = self.modules.discriminator(clean_wavs).mean(dim=1)
        simu_result = self.modules.discriminator(predictions).mean(dim=1)

        real_cost = 0
        simu_cost = 0
        map_cost = self.hparams.compute_cost(predictions, clean_wavs, lens)

        # One is real, zero is fake
        if optim_name == "generator":
            simu_target = torch.ones(batch_size, 1, device=self.device)
            simu_cost = self.hparams.compute_cost(simu_result, simu_target)
            real_cost = 0.0
            self.metrics["G"].append(simu_cost.detach())
        elif optim_name == "discriminator":
            real_target = torch.ones(batch_size, 1, device=self.device)
            simu_target = torch.zeros(batch_size, 1, device=self.device)
            real_cost = self.hparams.compute_cost(real_result, real_target)
            simu_cost = self.hparams.compute_cost(simu_result, simu_target)
            self.metrics["D"].append((real_cost + simu_cost).detach())

        return real_cost + simu_cost + map_cost

    def fit_batch(self, batch):
        "Trains the GAN with a batch"
        self.g_optimizer.zero_grad()
        predictions = self.compute_forward(batch, sb.Stage.TRAIN)
        g_loss = self.compute_objectives(
            predictions, batch, sb.Stage.TRAIN, "generator"
        )
        g_loss.backward()
        self.g_optimizer.step()

        self.d_optimizer.zero_grad()
        predictions = self.compute_forward(batch, sb.Stage.TRAIN)
        d_loss = self.compute_objectives(
            predictions, batch, sb.Stage.TRAIN, "discriminator"
        )
        d_loss.backward()
        self.d_optimizer.step()

        return g_loss.detach() + d_loss.detach()

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage == sb.Stage.TRAIN:
            self.metrics = {"G": [], "D": []}

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            g_loss = torch.tensor(self.metrics["G"])
            d_loss = torch.tensor(self.metrics["D"])
            print("Avg G loss: %.2f" % torch.mean(g_loss))
            print("Avg D loss: %.2f" % torch.mean(d_loss))
            print("train loss: ", stage_loss)
        elif stage == sb.Stage.VALID:
            print("Completed epoch %d" % epoch)
            print("Valid loss: %.3f" % stage_loss)
        else:
            self.test_loss = stage_loss

    def init_optimizers(self):
        """Initializes the generator and discriminator optimizers"""
        self.g_optimizer = self.hparams.g_opt_class(
            self.modules.generator.parameters()
        )
        self.d_optimizer = self.hparams.d_opt_class(
            self.modules.discriminator.parameters()
        )

    def zero_grad(self, set_to_none=False):
        """Sets the gradients of all optimized `torch.Tensor`s to zero."""
        self.g_optimizer.zero_grad(set_to_none)
        self.d_optimizer.zero_grad(set_to_none)


def data_prep(data_folder):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder)

    # Trainer initialization
    gan_brain = EnhanceGanBrain(
        modules=hparams["modules"], hparams=hparams, run_opts={"device": device}
    )

    # Training/validation loop
    gan_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    gan_brain.evaluate(valid_data)

    # Check test loss (mse), train loss is GAN loss
    assert gan_brain.test_loss < 0.002


if __name__ == "__main__":
    main()


def test_loss(device):
    main(device)

tests/integration/enhance_GAN/example_enhance_gan_experiment.py
TzuchengChang/NASS
"""This minimal example checks on sampling with ascending/descending sorting and random shuffling; w/ & w/o DDP.
"""

import os
import torch
import pickle
import pathlib
import itertools
import speechbrain as sb
import torch.multiprocessing as mp
from hyperpyyaml import load_hyperpyyaml


class SamplingBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the binary probability."
        batch = batch.to(self.device)
        lens = batch.duration

        if stage == sb.Stage.TRAIN:
            self.ids_list.append(batch.id)
            if self.hparams.sorting == "ascending":
                # ignore last; non-evenly divisible data; 99 items -> last batch: 19 -> 20 items (thus, nearby sort)
                if not all(
                    [x == y for x, y in zip(lens[:-1], sorted(lens[:-1]))]
                ):  # ":-1" is specific to dummy data
                    print(lens)
                    assert False
            elif self.hparams.sorting == "descending":
                if not all(
                    [
                        x == y
                        for x, y in zip(
                            lens[:-1], sorted(lens[:-1], reverse=True)
                        )
                    ]
                ):
                    print(lens)
                    assert False
            elif self.hparams.sorting == "random":
                assert not all(
                    [x == y for x, y in zip(lens[:-1], sorted(lens[:-1]))]
                )
            else:
                raise NotImplementedError(
                    "sorting must be random, ascending or descending"
                )

        return lens

    def compute_objectives(self, predictions, batch, stage=True):
        "Given the network predictions and targets computed the binary CE"
        inputs = torch.tensor([10.0, -6.0], requires_grad=True)
        targets = torch.tensor([1, 0])
        loss = self.hparams.compute_loss(inputs, targets)
        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage == sb.Stage.TRAIN:
            self.ids_list = []

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            # check that all IDs are unique; no duplicate IDs
            batched_ids = sorted(list(itertools.chain(*self.ids_list)))
            assert batched_ids == sorted(list(set(batched_ids)))

            # write out to check later all IDs were visited
            if self.distributed_launch:
                with open(
                    f"tests/tmp/ddp_sorting_ids_{self.hparams.sorting}_{self.hparams.rank}",
                    "wb",
                ) as f:
                    pickle.dump(batched_ids, f)


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    train_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=data_folder / "annotation/dev-clean.csv",
        replacements={"data_root": data_folder},
    )

    # start: sorting impact
    if hparams["sorting"] == "ascending":
        # we sort training data to speed up training and get better results.
        train_data = train_data.filtered_sorted(
            sort_key="duration",
            # key_max_value={"duration": hparams["avoid_if_longer_than"]},
        )
        # when sorting do not shuffle in dataloader ! otherwise is pointless
        hparams["dataloader_options"]["shuffle"] = False
        # hparams["dataloader_options"]["drop_last"] = True  # drops last entry which is out of order

    elif hparams["sorting"] == "descending":
        train_data = train_data.filtered_sorted(
            sort_key="duration",
            reverse=True,
            # key_max_value={"duration": hparams["avoid_if_longer_than"]},
        )
        # when sorting do not shuffle in dataloader ! otherwise is pointless
        hparams["dataloader_options"]["shuffle"] = False
        # hparams["dataloader_options"]["drop_last"] = True  # drops last entry which is out of order

    elif hparams["sorting"] == "random":
        # hparams["dataloader_options"]["drop_last"] = True  # reduced length from 99 to 80
        pass

    else:
        raise NotImplementedError(
            "sorting must be random, ascending or descending"
        )
    # end: sorting impact

    valid_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=data_folder / "annotation/dev-clean.csv",
        replacements={"data_root": data_folder},
    )

    datasets = [train_data, valid_data]

    # 1. Define audio pipeline:
    @sb.utils.data_pipeline.takes("duration")
    @sb.utils.data_pipeline.provides("duration")
    def audio_pipeline(duration):
        return duration

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "duration"])

    return train_data, valid_data


def recipe(device="cpu", yaml_file="hyperparams.yaml", run_opts=None):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = os.path.join(experiment_dir, yaml_file)
    data_folder = "../../samples/"
    data_folder = (experiment_dir / data_folder).resolve()
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # usually here: sb.utils.distributed.ddp_init_group(run_opts)

    # Data IO creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    if run_opts is None:
        run_opts = {}
    else:
        hparams["rank"] = run_opts["local_rank"]
    run_opts["device"] = device

    ctc_brain = SamplingBrain(
        hparams["modules"], hparams["opt_class"], hparams, run_opts=run_opts,
    )

    # Training/validation loop
    ctc_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )


if __name__ == "__main__":
    recipe(yaml_file="random.yaml")
    recipe(yaml_file="asc.yaml")
    recipe(yaml_file="dsc.yaml")


def test_error(device):
    recipe(device=device, yaml_file="random.yaml")
    recipe(device=device, yaml_file="asc.yaml")
    recipe(device=device, yaml_file="dsc.yaml")


def ddp_recipes(rank, size, backend="gloo"):
    """ Initialize the distributed environment. """
    os.environ["WORLD_SIZE"] = f"{size}"
    os.environ["RANK"] = f"{rank}"
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29500"

    run_opts = dict()
    run_opts["distributed_launch"] = True
    run_opts["distributed_backend"] = backend
    run_opts["local_rank"] = rank

    sb.utils.distributed.ddp_init_group(run_opts)

    recipe(device="cpu", yaml_file="random.yaml", run_opts=run_opts)
    recipe(device="cpu", yaml_file="asc.yaml", run_opts=run_opts)
    recipe(device="cpu", yaml_file="dsc.yaml", run_opts=run_opts)


def test_ddp():
    size = 2
    processes = []
    mp.set_start_method("spawn", force=True)
    os.makedirs("tests/tmp", exist_ok=True)
    for rank in range(size):
        p = mp.Process(target=ddp_recipes, args=(rank, size))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
        assert p.exitcode == 0

    # check all
    for sorting in ["random", "ascending", "descending"]:
        ids = []
        for rank in range(2):
            idf = f"tests/tmp/ddp_sorting_ids_{sorting}_{rank}"
            with open(idf, "rb") as f:
                ids += pickle.load(f)
        assert (
            len(ids) == 100 if sorting == "random" else 99
        )  # sorted data stays within the 99; random not
        assert len(set(ids)) == 99

tests/integration/sampling/example_sorting.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a speech separation system with on a tiny dataset.
The architecture is based on ConvTasnet and expects in input mixtures of two
speakers.
"""

import torch
import pathlib
import speechbrain as sb
import torch.nn.functional as F
from hyperpyyaml import load_hyperpyyaml
from speechbrain.nnet.losses import get_si_snr_with_pitwrapper


class SepBrain(sb.Brain):
    def compute_forward(self, mixture, stage):
        "Given an input batch it computes the two estimated sources."
        mixture = mixture.to(self.device)
        mix_w = self.hparams.encoder(mixture)
        est_mask = self.hparams.mask_net(mix_w)
        mix_w = torch.stack([mix_w] * 2)
        sep_h = mix_w * est_mask

        # Decoding
        est_source = torch.cat(
            [self.hparams.decoder(sep_h[i]).unsqueeze(-1) for i in range(2)],
            dim=-1,
        )

        # T changed after conv1d in encoder, fix it here
        T_origin = mixture.size(1)
        T_conv = est_source.size(1)
        if T_origin > T_conv:
            est_source = F.pad(est_source, (0, 0, 0, T_origin - T_conv))
        else:
            est_source = est_source[:, :T_origin, :]

        return est_source

    def compute_objectives(self, predictions, targets):
        "Given the network predictions and targets computed the PIT loss."
        loss = get_si_snr_with_pitwrapper(targets, predictions)
        return loss

    def fit_batch(self, batch):
        """Fits a training batch."""
        inputs = batch.mix_sig.data.to(self.device)
        targets = torch.cat(
            [
                batch.source1.data.unsqueeze(-1),
                batch.source2.data.unsqueeze(-1),
            ],
            dim=-1,
        ).to(self.device)

        predictions = self.compute_forward(inputs, sb.Stage.TRAIN)
        loss = self.compute_objectives(predictions, targets)

        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
        return loss.detach()

    def evaluate_batch(self, batch, stage):
        """Evaluates a batch"""
        inputs = batch.mix_sig.data.to(self.device)
        targets = torch.cat(
            [
                batch.source1.data.unsqueeze(-1),
                batch.source2.data.unsqueeze(-1),
            ],
            dim=-1,
        ).to(self.device)

        predictions = self.compute_forward(inputs, stage)
        loss = self.compute_objectives(predictions, targets)
        return loss.detach()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        elif stage == sb.Stage.VALID:
            print("Completed epoch %d" % epoch)
            print("Train SI-SNR: %.3f" % -self.train_loss)
            print("Valid SI-SNR: %.3f" % -stage_loss)
        elif stage == sb.Stage.TEST:
            print("Test SI-SNR: %.3f" % -stage_loss)


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=data_folder / "../annotation/separation_train.csv",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=data_folder / "../annotation/separation_dev.csv",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("mix_wav", "s1_wav", "s2_wav")
    @sb.utils.data_pipeline.provides("mix_sig", "source1", "source2")
    def audio_pipeline(mix_wav, s1_wav, s2_wav):
        mix_sig = sb.dataio.dataio.read_audio(mix_wav)
        yield mix_sig
        source1 = sb.dataio.dataio.read_audio(s1_wav)
        yield source1
        source2 = sb.dataio.dataio.read_audio(s2_wav)
        yield source2

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets, ["id", "mix_sig", "source1", "source2"]
    )

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/separation"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    sep_brain = SepBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    sep_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    sep_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert sep_brain.train_loss < 5.0


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/separation/example_conv_tasnet.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This minimal example trains a speaker identification system based on
x-vectors. The encoder is based on TDNNs. The classifier is a MLP.
"""

import pathlib
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml


# Trains xvector model
class XvectorBrain(sb.Brain):
    def compute_forward(self, batch, stage):
        "Given an input batch it computes the speaker probabilities."
        batch = batch.to(self.device)
        wavs, lens = batch.sig
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.mean_var_norm(feats, lens)
        x_vect = self.modules.xvector_model(feats)
        outputs = self.modules.classifier(x_vect)

        return outputs, lens

    def compute_objectives(self, predictions, batch, stage):
        "Given the network predictions and targets computed the CE loss."
        predictions, lens = predictions
        spkid, spkid_lens = batch.spk_id_encoded
        loss = self.hparams.compute_cost(predictions, spkid, lens)

        if stage != sb.Stage.TRAIN:
            self.error_metrics.append(batch.id, predictions, spkid, lens)

        return loss

    def on_stage_start(self, stage, epoch=None):
        "Gets called when a stage (either training, validation, test) starts."
        if stage != sb.Stage.TRAIN:
            self.error_metrics = self.hparams.error_stats()

    def on_stage_end(self, stage, stage_loss, epoch=None):
        """Gets called at the end of a stage."""
        if stage == sb.Stage.TRAIN:
            self.train_loss = stage_loss
        if stage == sb.Stage.VALID:
            print("Epoch %d complete" % epoch)
            print("Train loss: %.2f" % self.train_loss)
        if stage != sb.Stage.TRAIN:
            print(stage, "loss: %.2f" % stage_loss)
            print(
                stage, "error: %.2f" % self.error_metrics.summarize("average")
            )


def data_prep(data_folder, hparams):
    "Creates the datasets and their data processing pipelines."

    # 1. Declarations:
    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_train.json",
        replacements={"data_root": data_folder},
    )
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(
        json_path=data_folder / "../annotation/ASR_dev.json",
        replacements={"data_root": data_folder},
    )
    datasets = [train_data, valid_data]
    label_encoder = sb.dataio.encoder.CategoricalEncoder()
    label_encoder.expect_len(hparams["num_spks"])

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("spk_id")
    @sb.utils.data_pipeline.provides("spk_id", "spk_id_encoded")
    def label_pipeline(spk_id):
        yield spk_id
        spk_id_encoded = label_encoder.encode_sequence_torch([spk_id])
        yield spk_id_encoded

    sb.dataio.dataset.add_dynamic_item(datasets, label_pipeline)

    # 3. Fit encoder:
    # NOTE: In this minimal example, also update from valid data
    label_encoder.update_from_didataset(train_data, output_key="spk_id")
    label_encoder.update_from_didataset(valid_data, output_key="spk_id")

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(datasets, ["id", "sig", "spk_id_encoded"])

    return train_data, valid_data


def main(device="cpu"):
    experiment_dir = pathlib.Path(__file__).resolve().parent
    hparams_file = experiment_dir / "hyperparams.yaml"
    data_folder = "../../samples/ASR"
    data_folder = (experiment_dir / data_folder).resolve()

    # Load model hyper parameters:
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin)

    # Dataset creation
    train_data, valid_data = data_prep(data_folder, hparams)

    # Trainer initialization
    xvect_brain = XvectorBrain(
        hparams["modules"],
        hparams["opt_class"],
        hparams,
        run_opts={"device": device},
    )

    # Training/validation loop
    xvect_brain.fit(
        range(hparams["N_epochs"]),
        train_data,
        valid_data,
        train_loader_kwargs=hparams["dataloader_options"],
        valid_loader_kwargs=hparams["dataloader_options"],
    )
    # Evaluation is run separately (now just evaluating on valid data)
    xvect_brain.evaluate(valid_data)

    # Check if model overfits for integration test
    assert xvect_brain.train_loss < 0.2


if __name__ == "__main__":
    main()


def test_error(device):
    main(device)

tests/integration/speaker_id/example_xvector_experiment.py
TzuchengChang/NASS
../../../templates/speech_recognition/ASR/train.py
tests/templates/fetching_ddp_dynbatch_finetuning/ASR_template_train.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""Recipe for debugging/testing, based upon SpeechBrain's minilibrispeech ASR template.

Does the following feature set work out together on some environment?
    DDP; dynamic batching; fine-tuning; mixed pretrainer fetching & testing using pretrained interface

Authors:
    * Andreas Nautsch 2023
"""
import os
import sys
import torch
import logging
import speechbrain as sb
from copy import deepcopy
from tqdm.contrib import tqdm
from torch.utils.data import DataLoader
from hyperpyyaml import load_hyperpyyaml
from speechbrain.pretrained import EncoderDecoderASR
from speechbrain.utils.distributed import run_on_main, ddp_barrier
from speechbrain.utils.data_utils import batch_pad_right
from speechbrain.dataio.dataset import DynamicItemDataset
from speechbrain.dataio.sampler import DynamicBatchSampler
from speechbrain.dataio.dataloader import (
    LoopedLoader,
    make_dataloader,
    distributed_loader_specifics,
)
from speechbrain.dataio.dataio import (
    read_audio,
    # read_audio_multichannel,
)
from ASR_template_train import ASR, dataio_prepare


logger = logging.getLogger(__name__)


def eval_reporting(reports, single_node=False):
    """Performance logging independent of who logs what.

    Parameters
    ----------
    reports: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    single_node: bool
        Flag for whether/not DDP-gather results (Default: False).
    """
    for log_metric, specs in reports.items():
        if not single_node:
            print(
                f'{log_metric} on DDP rank {int(os.environ["RANK"])}: {specs["tracker"].summarize()}'
            )
            result_list = [None for _ in range(int(os.environ["WORLD_SIZE"]))]
            # WARNING: https://pytorch.org/docs/stable/distributed.html - underlying `pickle` module is known to be insecure
            torch.distributed.all_gather_object(
                result_list, specs["tracker"].scores
            )
            specs["tracker"].scores = list()
            for r in result_list:
                specs["tracker"].scores.extend(r)
        summary = specs["tracker"].summarize()
        print(f"\tSummary: {summary}")
        logger.info(f'{log_metric}: {summary[specs["field"]]}\n')


def eval_test_use_recipe_dataio(
    encoder_decoder_asr, test_set, test_kwargs, reporter, single_node=False
):
    """Bypassing speechbrain.pretrained.Pretrained.load_audio with recipe dataio (speechbrain.dataio.dataio.read_audio).

    Parameters
    ----------
    encoder_decoder_asr: speechbrain.pretrained.EncoderDecoderASR
        Pretrained interface (other interfaces will require other functions to be called; this is an example).
    test_set: dict
        Data loader options for testing.
    test_kwargs: dict
        Data loader options for testing.
    reporter: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    single_node: bool
        Flag for whether/not DDP-gather results (Default: False).
    """
    if "ckpt_prefix" in test_kwargs:
        del test_kwargs["ckpt_prefix"]
    if not (
        isinstance(test_set, DataLoader) or isinstance(test_set, LoopedLoader)
    ):
        test_set = make_dataloader(test_set, **test_kwargs)

    with torch.no_grad():
        for batch in tqdm(test_set, dynamic_ncols=True, disable=False):
            batch = batch.to(encoder_decoder_asr.device)
            predictions = encoder_decoder_asr.transcribe_batch(*batch.sig)

            # prepare for metric reporting
            predicted = [wrd.split(" ") for wrd in predictions[0]]
            targeted = [wrd.split(" ") for wrd in batch.words]
            ids = batch.id
            for metric in reporter.keys():
                reporter[metric]["tracker"].append(
                    ids=ids, predict=predicted, target=targeted
                )

        # Report summary
        eval_reporting(reports=reporter, single_node=single_node)


def eval_test_batch_from_scratch(
    encoder_decoder_asr,
    test_set,
    test_kwargs,
    reporter,
    pretrainer_load_audio=False,
):
    """Relies only on batched audio paths to create batches using the pretrained interface only.

    Parameters
    ----------
    encoder_decoder_asr: speechbrain.pretrained.EncoderDecoderASR
        Pretrained interface (other interfaces will require other functions to be called; this is an example).
    test_set: Dataset, DataLoader
        If a DataLoader is given, it is iterated directly. Otherwise passed to `sb.dataio.dataloader.make_dataloader()`.
    test_kwargs: dict
        Data loader options for testing.
    reporter: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    pretrainer_load_audio: bool (Default: False)
        Whether to use Pretrainer.load_audio (True) or dataio.read_audio (False).
    """
    if "ckpt_prefix" in test_kwargs:
        del test_kwargs["ckpt_prefix"]
    if not (
        isinstance(test_set, DataLoader) or isinstance(test_set, LoopedLoader)
    ):
        test_set = make_dataloader(test_set, **test_kwargs)

    with torch.no_grad():
        for batch in tqdm(test_set, dynamic_ncols=True, disable=False):
            # instead of using batch.sig, we desire to see pretrained_hf_asr.load_audio in action
            wavs = []
            for audio_path in batch.wav:  # get the paths only
                if pretrainer_load_audio:
                    wavs.append(
                        encoder_decoder_asr.load_audio(
                            path=audio_path, silent_local_fetch=True
                        )
                    )
                else:
                    wavs.append(
                        read_audio(audio_path).to(encoder_decoder_asr.device)
                    )
            wavs, wav_lens = batch_pad_right(wavs)
            predictions = encoder_decoder_asr(wavs, wav_lens)

            # prepare for metric reporting
            predicted = [wrd.split(" ") for wrd in predictions[0]]
            targeted = [wrd.split(" ") for wrd in batch.words]
            ids = batch.id
            for metric in reporter.keys():
                reporter[metric]["tracker"].append(
                    ids=ids, predict=predicted, target=targeted
                )

        # Report summary
        ddp_barrier()
        eval_reporting(reports=reporter)


if __name__ == "__main__":
    # CLI parsing
    hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])

    # DDP init
    sb.utils.distributed.ddp_init_group(run_opts)

    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    # Kept aside for later
    reporting = {
        "WER": {
            "field": "error_rate",
            "tracker": deepcopy(
                hparams["error_rate_computer"]()
            ),  # n_jobs=int(os.environ['WORLD_SIZE']))),
        },
        "CER": {
            "field": "error_rate",
            "tracker": deepcopy(
                hparams["cer_computer"]()
            ),  # n_jobs=int(os.environ['WORLD_SIZE']))),
        },
    }

    # Create experiment directory
    sb.create_experiment_directory(
        experiment_directory=hparams["output_folder"],
        hyperparams_to_save=hparams_file,
        overrides=overrides,
    )

    # Spared: dataset prepare script & its run_on_main - prepare dataio pipeline only
    datasets = dataio_prepare(hparams)

    # Add: dynamic batching (only in training & validation)
    dynamic_hparams = hparams["dynamic_batch_sampler"]
    train_dataloader_opts = hparams["train_dataloader_opts"]
    valid_dataloader_opts = hparams["valid_dataloader_opts"]
    train_dataloader_opts = {
        "batch_sampler": DynamicBatchSampler(
            datasets["train"],
            dynamic_hparams["max_batch_len"],
            num_buckets=dynamic_hparams["num_buckets"],
            length_func=lambda x: x["length"],
            shuffle=dynamic_hparams["shuffle_ex"],
            batch_ordering=dynamic_hparams["batch_ordering"],
        ),
        "num_workers": train_dataloader_opts["num_workers"],
    }
    valid_dataloader_opts = {
        "batch_sampler": DynamicBatchSampler(
            datasets["valid"],
            dynamic_hparams["max_batch_len"],
            num_buckets=dynamic_hparams["num_buckets"],
            length_func=lambda x: x["length"],
            shuffle=dynamic_hparams["shuffle_ex"],
            batch_ordering=dynamic_hparams["batch_ordering"],
        ),
        "num_workers": valid_dataloader_opts["num_workers"],
    }

    # We download:
    # * the tokenizer from URL - https://huggingface.co/speechbrain/asr-crdnn-rnnlm-librispeech/resolve/main/...
    # * the pretrained LM from HuggingFace - HF: speechbrain/asr-crdnn-rnnlm-librispeech
    # * the pretrained ASR from the local template checkpoint - local: speechbrain/asr-crdnn-rnnlm-librispeech
    run_on_main(hparams["pretrainer_tokenizer"].collect_files,)
    run_on_main(hparams["pretrainer_LM"].collect_files,)
    hparams["pretrainer_tokenizer"].load_collected(run_opts["device"])
    hparams["pretrainer_LM"].load_collected(run_opts["device"])
    # LOCAL fetching takes sources directly from their location
    hparams["pretrainer_ASR"].collect_files()
    hparams["pretrainer_ASR"].load_collected(run_opts["device"])

    # Trainer initialization
    asr_brain = ASR(
        modules=hparams["modules"],
        opt_class=hparams["opt_class"],
        hparams=hparams,
        run_opts=deepcopy(run_opts),
        checkpointer=hparams["checkpointer"],
    )
    asr_brain.hparams.valid_search = asr_brain.hparams.valid_search.to(
        asr_brain.device
    )
    asr_brain.hparams.test_search = asr_brain.hparams.test_search.to(
        asr_brain.device
    )
    asr_brain.cer_metric = deepcopy(reporting["CER"]["tracker"])
    asr_brain.wer_metric = deepcopy(reporting["WER"]["tracker"])

    # Freeze all but LM
    for mod in [
        "encoder",
        "embedding",
    ]:  # decoder, ctc_lin & seq_lin are for fine-tuning
        for param in getattr(asr_brain.modules, mod).parameters():
            param.requires_grad = False

    # Fine-tuning
    asr_brain.fit(
        asr_brain.hparams.epoch_counter,
        datasets["train"],
        datasets["valid"],
        train_loader_kwargs=train_dataloader_opts,
        valid_loader_kwargs=valid_dataloader_opts,
    )

    # Testing
    asr_brain.evaluate(
        test_set=datasets["test"],
        min_key="WER",
        test_loader_kwargs=hparams["test_dataloader_opts"],
    )

    # Save, so it can be found as pre-trained model source
    if not os.path.exists(f"{hparams['save_folder']}/CKPT+latest"):
        ddp_barrier()  # just to be sure
        run_on_main(
            asr_brain.checkpointer.save_checkpoint, kwargs={"name": "latest"}
        )

    # Clean-up memory (for using EncoderDecoderASR interface only) but preserve testing-relevant objects
    test_datasets = deepcopy(datasets["test"])
    test_loader_kwargs = deepcopy(hparams["test_dataloader_opts"])
    del asr_brain, datasets

    # Pre-trained interface
    pretrained_asr = EncoderDecoderASR.from_hparams(
        source="source_pretrained",
        savedir="source_pretrained",
        hparams_file="pretrained.yaml",
        run_opts=deepcopy(run_opts),
    )

    # Test w/ DDP
    ddp_test_set = DynamicItemDataset.from_json(
        json_path=hparams["test_annotation"],
        replacements={"data_root": hparams["data_folder"]},
        dynamic_items=[],
        output_keys=["id", "wav", "words"],
    )
    ddp_test_kwargs = distributed_loader_specifics(
        distributed_launch=True,
        rank=int(os.environ["RANK"]),
        dataset=ddp_test_set,
        loader_kwargs=deepcopy(test_loader_kwargs),
    )
    for flag in [True, False]:
        logger.info(f"\nBatch from scratch w/ pretrainer_load_audio={flag}")
        eval_test_batch_from_scratch(
            encoder_decoder_asr=pretrained_asr,
            test_set=deepcopy(ddp_test_set),
            test_kwargs=deepcopy(ddp_test_kwargs),
            reporter=deepcopy(reporting),
            pretrainer_load_audio=flag,
        )

    # Re:testing w/ previous dataloader // note: needs to run as last item (the script might get stuck otherwise)
    logger.info(f"\nTesting w/ asr_brain's eval dataloader")
    run_on_main(
        eval_test_use_recipe_dataio,
        kwargs={
            "encoder_decoder_asr": pretrained_asr,
            "test_set": deepcopy(test_datasets),
            "test_kwargs": deepcopy(test_loader_kwargs),
            "reporter": deepcopy(reporting),
            "single_node": True,
        },
    )

tests/templates/fetching_ddp_dynbatch_finetuning/finetune.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""Recipe for debugging/testing, based upon SpeechBrain's minilibrispeech ASR template.

Does the following feature set work out together on some environment?
    DDP; dynamic batching; fine-tuning; mixed pretrainer fetching & testing using pretrained interface

Authors:
    * Andreas Nautsch 2023
"""
import os
import sys
import torch
import logging
import speechbrain as sb
from copy import deepcopy
from tqdm.contrib import tqdm
from torch.utils.data import DataLoader
from hyperpyyaml import load_hyperpyyaml
from speechbrain.pretrained import EncoderDecoderASR
from speechbrain.utils.distributed import run_on_main, ddp_barrier
from speechbrain.utils.data_utils import batch_pad_right
from speechbrain.dataio.dataset import DynamicItemDataset
from speechbrain.dataio.sampler import DynamicBatchSampler
from speechbrain.dataio.dataloader import (
    LoopedLoader,
    make_dataloader,
    distributed_loader_specifics,
)
from speechbrain.dataio.dataio import (
    read_audio,
    # read_audio_multichannel,
)
from ASR_template_train import ASR, dataio_prepare


logger = logging.getLogger(__name__)


def eval_reporting(reports, single_node=False):
    """Performance logging independent of who logs what.

    Parameters
    ----------
    reports: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    single_node: bool
        Flag for whether/not DDP-gather results (Default: False).
    """
    for log_metric, specs in reports.items():
        if not single_node:
            print(
                f'{log_metric} on DDP rank {int(os.environ["RANK"])}: {specs["tracker"].summarize()}'
            )
            result_list = [None for _ in range(int(os.environ["WORLD_SIZE"]))]
            # WARNING: https://pytorch.org/docs/stable/distributed.html - underlying `pickle` module is known to be insecure
            torch.distributed.all_gather_object(
                result_list, specs["tracker"].scores
            )
            specs["tracker"].scores = list()
            for r in result_list:
                specs["tracker"].scores.extend(r)
        summary = specs["tracker"].summarize()
        print(f"\tSummary: {summary}")
        logger.info(f'{log_metric}: {summary[specs["field"]]}\n')


def eval_test_use_recipe_dataio(
    encoder_decoder_asr, test_set, test_kwargs, reporter, single_node=False
):
    """Bypassing speechbrain.pretrained.Pretrained.load_audio with recipe dataio (speechbrain.dataio.dataio.read_audio).

    Parameters
    ----------
    encoder_decoder_asr: speechbrain.pretrained.EncoderDecoderASR
        Pretrained interface (other interfaces will require other functions to be called; this is an example).
    test_set: dict
        Data loader options for testing.
    test_kwargs: dict
        Data loader options for testing.
    reporter: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    single_node: bool
        Flag for whether/not DDP-gather results (Default: False).
    """
    if "ckpt_prefix" in test_kwargs:
        del test_kwargs["ckpt_prefix"]
    if not (
        isinstance(test_set, DataLoader) or isinstance(test_set, LoopedLoader)
    ):
        test_set = make_dataloader(test_set, **test_kwargs)

    with torch.no_grad():
        for batch in tqdm(test_set, dynamic_ncols=True, disable=False):
            batch = batch.to(encoder_decoder_asr.device)
            predictions = encoder_decoder_asr.transcribe_batch(*batch.sig)

            # prepare for metric reporting
            predicted = [wrd.split(" ") for wrd in predictions[0]]
            targeted = [wrd.split(" ") for wrd in batch.words]
            ids = batch.id
            for metric in reporter.keys():
                reporter[metric]["tracker"].append(
                    ids=ids, predict=predicted, target=targeted
                )

        # Report summary
        eval_reporting(reports=reporter, single_node=single_node)


def eval_test_batch_from_scratch(
    encoder_decoder_asr,
    test_set,
    test_kwargs,
    reporter,
    pretrainer_load_audio=False,
):
    """Relies only on batched audio paths to create batches using the pretrained interface only.

    Parameters
    ----------
    encoder_decoder_asr: speechbrain.pretrained.EncoderDecoderASR
        Pretrained interface (other interfaces will require other functions to be called; this is an example).
    test_set: Dataset, DataLoader
        If a DataLoader is given, it is iterated directly. Otherwise passed to `sb.dataio.dataloader.make_dataloader()`.
    test_kwargs: dict
        Data loader options for testing.
    reporter: dict
        Maps metric labels to performance trackers (instances) which need summarise certain fields for final reporting.
    pretrainer_load_audio: bool (Default: False)
        Whether to use Pretrainer.load_audio (True) or dataio.read_audio (False).
    """
    if "ckpt_prefix" in test_kwargs:
        del test_kwargs["ckpt_prefix"]
    if not (
        isinstance(test_set, DataLoader) or isinstance(test_set, LoopedLoader)
    ):
        test_set = make_dataloader(test_set, **test_kwargs)

    with torch.no_grad():
        for batch in tqdm(test_set, dynamic_ncols=True, disable=False):
            # instead of using batch.sig, we desire to see pretrained_hf_asr.load_audio in action
            wavs = []
            for audio_path in batch.wav:  # get the paths only
                if pretrainer_load_audio:
                    wavs.append(
                        encoder_decoder_asr.load_audio(
                            path=audio_path, silent_local_fetch=True
                        )
                    )
                else:
                    wavs.append(
                        read_audio(audio_path).to(encoder_decoder_asr.device)
                    )
            wavs, wav_lens = batch_pad_right(wavs)
            predictions = encoder_decoder_asr(wavs, wav_lens)

            # prepare for metric reporting
            predicted = [wrd.split(" ") for wrd in predictions[0]]
            targeted = [wrd.split(" ") for wrd in batch.words]
            ids = batch.id
            for metric in reporter.keys():
                reporter[metric]["tracker"].append(
                    ids=ids, predict=predicted, target=targeted
                )

        # Report summary
        ddp_barrier()
        eval_reporting(reports=reporter)


if __name__ == "__main__":
    # CLI parsing
    hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])

    # DDP init
    sb.utils.distributed.ddp_init_group(run_opts)

    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    # Kept aside for later
    reporting = {
        "WER": {
            "field": "error_rate",
            "tracker": deepcopy(
                hparams["error_rate_computer"]()
            ),  # n_jobs=int(os.environ['WORLD_SIZE']))),
        },
        "CER": {
            "field": "error_rate",
            "tracker": deepcopy(
                hparams["cer_computer"]()
            ),  # n_jobs=int(os.environ['WORLD_SIZE']))),
        },
    }

    # Create experiment directory
    sb.create_experiment_directory(
        experiment_directory=hparams["output_folder"],
        hyperparams_to_save=hparams_file,
        overrides=overrides,
    )

    # Spared: dataset prepare script & its run_on_main - prepare dataio pipeline only
    datasets = dataio_prepare(hparams)

    # Add: dynamic batching (only in training & validation)
    dynamic_hparams = hparams["dynamic_batch_sampler"]
    train_dataloader_opts = hparams["train_dataloader_opts"]
    valid_dataloader_opts = hparams["valid_dataloader_opts"]
    train_dataloader_opts = {
        "batch_sampler": DynamicBatchSampler(
            datasets["train"],
            dynamic_hparams["max_batch_len"],
            num_buckets=dynamic_hparams["num_buckets"],
            length_func=lambda x: x["length"],
            shuffle=dynamic_hparams["shuffle_ex"],
            batch_ordering=dynamic_hparams["batch_ordering"],
        ),
        "num_workers": train_dataloader_opts["num_workers"],
    }
    valid_dataloader_opts = {
        "batch_sampler": DynamicBatchSampler(
            datasets["valid"],
            dynamic_hparams["max_batch_len"],
            num_buckets=dynamic_hparams["num_buckets"],
            length_func=lambda x: x["length"],
            shuffle=dynamic_hparams["shuffle_ex"],
            batch_ordering=dynamic_hparams["batch_ordering"],
        ),
        "num_workers": valid_dataloader_opts["num_workers"],
    }

    # We download:
    # * the tokenizer from URL - https://huggingface.co/speechbrain/asr-crdnn-rnnlm-librispeech/resolve/main/...
    # * the pretrained LM from HuggingFace - HF: speechbrain/asr-crdnn-rnnlm-librispeech
    # * the pretrained ASR from the local template checkpoint - local: speechbrain/asr-crdnn-rnnlm-librispeech

    # === this is the only block that changed to finetune.py
    hparams["pretrainer"].collect_files(internal_ddp_handling=True)
    hparams["pretrainer"].load_collected(run_opts["device"])
    # ===

    # Trainer initialization
    asr_brain = ASR(
        modules=hparams["modules"],
        opt_class=hparams["opt_class"],
        hparams=hparams,
        run_opts=deepcopy(run_opts),
        checkpointer=hparams["checkpointer"],
    )
    asr_brain.hparams.valid_search = asr_brain.hparams.valid_search.to(
        asr_brain.device
    )
    asr_brain.hparams.test_search = asr_brain.hparams.test_search.to(
        asr_brain.device
    )
    asr_brain.cer_metric = deepcopy(reporting["CER"]["tracker"])
    asr_brain.wer_metric = deepcopy(reporting["WER"]["tracker"])

    # Freeze all but LM
    for mod in [
        "encoder",
        "embedding",
    ]:  # decoder, ctc_lin & seq_lin are for fine-tuning
        for param in getattr(asr_brain.modules, mod).parameters():
            param.requires_grad = False

    # Fine-tuning
    asr_brain.fit(
        asr_brain.hparams.epoch_counter,
        datasets["train"],
        datasets["valid"],
        train_loader_kwargs=train_dataloader_opts,
        valid_loader_kwargs=valid_dataloader_opts,
    )

    # Testing
    asr_brain.evaluate(
        test_set=datasets["test"],
        min_key="WER",
        test_loader_kwargs=hparams["test_dataloader_opts"],
    )

    # Save, so it can be found as pre-trained model source
    if not os.path.exists(f"{hparams['save_folder']}/CKPT+latest"):
        ddp_barrier()  # just to be sure
        run_on_main(
            asr_brain.checkpointer.save_checkpoint, kwargs={"name": "latest"}
        )

    # Clean-up memory (for using EncoderDecoderASR interface only) but preserve testing-relevant objects
    test_datasets = deepcopy(datasets["test"])
    test_loader_kwargs = deepcopy(hparams["test_dataloader_opts"])
    del asr_brain, datasets

    # Pre-trained interface
    pretrained_asr = EncoderDecoderASR.from_hparams(
        source="source_pretrained",
        savedir="source_pretrained",
        hparams_file="pretrained.yaml",
        run_opts=deepcopy(run_opts),
    )

    # Test w/ DDP
    ddp_test_set = DynamicItemDataset.from_json(
        json_path=hparams["test_annotation"],
        replacements={"data_root": hparams["data_folder"]},
        dynamic_items=[],
        output_keys=["id", "wav", "words"],
    )
    ddp_test_kwargs = distributed_loader_specifics(
        distributed_launch=True,
        rank=int(os.environ["RANK"]),
        dataset=ddp_test_set,
        loader_kwargs=deepcopy(test_loader_kwargs),
    )
    for flag in [True, False]:
        logger.info(f"\nBatch from scratch w/ pretrainer_load_audio={flag}")
        eval_test_batch_from_scratch(
            encoder_decoder_asr=pretrained_asr,
            test_set=deepcopy(ddp_test_set),
            test_kwargs=deepcopy(ddp_test_kwargs),
            reporter=deepcopy(reporting),
            pretrainer_load_audio=flag,
        )

    # Re:testing w/ previous dataloader // note: needs to run as last item (the script might get stuck otherwise)
    logger.info(f"\nTesting w/ asr_brain's eval dataloader")
    run_on_main(
        eval_test_use_recipe_dataio,
        kwargs={
            "encoder_decoder_asr": pretrained_asr,
            "test_set": deepcopy(test_datasets),
            "test_kwargs": deepcopy(test_loader_kwargs),
            "reporter": deepcopy(reporting),
            "single_node": True,
        },
    )

tests/templates/fetching_ddp_dynbatch_finetuning/finetune_fetch_once.py
TzuchengChang/NASS
../../../templates/speech_recognition/ASR/mini_librispeech_prepare.py
tests/templates/fetching_ddp_dynbatch_finetuning/mini_librispeech_prepare.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""Recipe for regression testing with an exemplary recipe; based on:
    recipes/timers-and-such/decoupled/hparams/train_TAS_LM.yaml
    by Loren Lugosch, Mirco Ravanelli 2020

Authors:
    * Andreas Nautsch 2023
"""

import sys
import torch
import speechbrain as sb
from hyperpyyaml import load_hyperpyyaml
from speechbrain.utils.distributed import run_on_main


# Define training procedure
class SLU(sb.Brain):
    """see recipes/timers-and-such/decoupled/train.py
    """

    def compute_forward(self, batch, stage):
        """Forward computations from the waveform batches to the output probabilities."""
        batch = batch.to(self.device)
        if stage == sb.Stage.TRAIN or stage == sb.Stage.VALID:
            # true_transcripts, true_transcript_lens = batch.transcript
            true_transcripts = batch.transcript
        else:
            wavs, wav_lens = batch.sig
        tokens_bos, tokens_bos_lens = batch.tokens_bos

        # Forward pass
        if stage == sb.Stage.TRAIN or stage == sb.Stage.VALID:
            asr_tokens = [
                self.hparams.asr_model.tokenizer.encode_as_ids(t)
                for t in true_transcripts
            ]
        else:
            words, asr_tokens = self.hparams.asr_model.transcribe_batch(
                wavs.detach(), wav_lens
            )

        # Pad examples to have same length.
        asr_tokens_lens = torch.tensor([max(len(t), 1) for t in asr_tokens])
        max_length = asr_tokens_lens.max().item()
        for t in asr_tokens:
            t += [0] * (max_length - len(t))
        asr_tokens = torch.tensor([t for t in asr_tokens])

        asr_tokens_lens = asr_tokens_lens.float()
        asr_tokens_lens = asr_tokens_lens / asr_tokens_lens.max()

        asr_tokens, asr_tokens_lens = (
            asr_tokens.to(self.device),
            asr_tokens_lens.to(self.device),
        )
        embedded_transcripts = self.hparams.input_emb(asr_tokens)

        # SLU forward pass
        encoder_out = self.hparams.slu_enc(embedded_transcripts)
        e_in = self.hparams.output_emb(tokens_bos)
        h, _ = self.hparams.dec(e_in, encoder_out, asr_tokens_lens)

        # Output layer for seq2seq log-probabilities
        logits = self.hparams.seq_lin(h)
        p_seq = self.hparams.log_softmax(logits)

        # Compute outputs
        if (
            stage == sb.Stage.TRAIN
            and self.batch_count % show_results_every != 0
        ):
            return p_seq, asr_tokens_lens
        else:
            p_tokens, scores = self.hparams.beam_searcher(
                encoder_out, asr_tokens_lens
            )
            return p_seq, asr_tokens_lens, p_tokens

    def compute_objectives(self, predictions, batch, stage):
        """Computes the loss (NLL) given predictions and targets."""

        if (
            stage == sb.Stage.TRAIN
            and self.batch_count % show_results_every != 0
        ):
            p_seq, asr_tokens_lens = predictions
        else:
            p_seq, asr_tokens_lens, predicted_tokens = predictions

        ids = batch.id
        tokens_eos, tokens_eos_lens = batch.tokens_eos
        tokens, tokens_lens = batch.tokens

        loss_seq = self.hparams.seq_cost(
            p_seq, tokens_eos, length=tokens_eos_lens
        )

        # (No ctc loss)
        loss = loss_seq

        if (stage != sb.Stage.TRAIN) or (
            self.batch_count % show_results_every == 0
        ):
            # Decode token terms to words
            predicted_semantics = [
                tokenizer.decode_ids(utt_seq).split(" ")
                for utt_seq in predicted_tokens
            ]

            target_semantics = [wrd.split(" ") for wrd in batch.semantics]

            for i in range(len(target_semantics)):
                print(" ".join(predicted_semantics[i]))
                print(" ".join(target_semantics[i]))
                print("")

            if stage != sb.Stage.TRAIN:
                self.wer_metric.append(
                    ids, predicted_semantics, target_semantics
                )
                self.cer_metric.append(
                    ids, predicted_semantics, target_semantics
                )

        return loss

    def fit_batch(self, batch):
        """Train the parameters given a single batch in input"""
        predictions = self.compute_forward(batch, sb.Stage.TRAIN)
        loss = self.compute_objectives(predictions, batch, sb.Stage.TRAIN)
        loss.backward()
        if self.check_gradients(loss):
            self.optimizer.step()
        self.optimizer.zero_grad()
        self.batch_count += 1
        return loss.detach()

    def evaluate_batch(self, batch, stage):
        """Computations needed for validation/test batches"""
        predictions = self.compute_forward(batch, stage=stage)
        loss = self.compute_objectives(predictions, batch, stage=stage)
        return loss.detach()

    def on_stage_start(self, stage, epoch):
        """Gets called at the beginning of each epoch"""
        self.batch_count = 0

        if stage != sb.Stage.TRAIN:

            self.cer_metric = self.hparams.cer_computer()
            self.wer_metric = self.hparams.error_rate_computer()

    def on_stage_end(self, stage, stage_loss, epoch):
        """Gets called at the end of a epoch."""
        # Compute/store important stats
        stage_stats = {"loss": stage_loss}
        if stage == sb.Stage.TRAIN:
            self.train_stats = stage_stats
        else:
            stage_stats["CER"] = self.cer_metric.summarize("error_rate")
            stage_stats["WER"] = self.wer_metric.summarize("error_rate")
            stage_stats["SER"] = self.wer_metric.summarize("SER")

        # Perform end-of-iteration things, like annealing, logging, etc.
        if stage == sb.Stage.VALID:
            old_lr, new_lr = self.hparams.lr_annealing(stage_stats["SER"])
            sb.nnet.schedulers.update_learning_rate(self.optimizer, new_lr)
            self.hparams.train_logger.log_stats(
                stats_meta={"epoch": epoch, "lr": old_lr},
                train_stats=self.train_stats,
                valid_stats=stage_stats,
            )
            self.checkpointer.save_and_keep_only(
                meta={"SER": stage_stats["SER"]}, min_keys=["SER"],
            )
        elif stage == sb.Stage.TEST:
            self.hparams.train_logger.log_stats(
                stats_meta={"Epoch loaded": self.hparams.epoch_counter.current},
                test_stats=stage_stats,
            )
            with open(self.hparams.wer_file, "w") as w:
                self.wer_metric.write_stats(w)


def data_io_prepare(hparams):
    """This function prepares the datasets to be used in the brain class.
    It also defines the data processing pipeline through user-defined functions."""

    data_folder = hparams["data_folder"]

    train_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["csv_train"], replacements={"data_root": data_folder},
    )

    if hparams["sorting"] == "ascending":
        # we sort training data to speed up training and get better results.
        train_data = train_data.filtered_sorted(sort_key="duration")
        # when sorting do not shuffle in dataloader ! otherwise is pointless
        hparams["dataloader_opts"]["shuffle"] = False

    elif hparams["sorting"] == "descending":
        train_data = train_data.filtered_sorted(
            sort_key="duration", reverse=True
        )
        # when sorting do not shuffle in dataloader ! otherwise is pointless
        hparams["dataloader_opts"]["shuffle"] = False

    elif hparams["sorting"] == "random":
        pass

    else:
        raise NotImplementedError(
            "sorting must be random, ascending or descending"
        )

    # If we are testing on all the real data, including dev-real,
    # we shouldn't use dev-real as the validation set.
    if hparams["test_on_all_real"]:
        valid_path = hparams["csv_dev_synth"]
    else:
        valid_path = hparams["csv_dev_real"]
    valid_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=valid_path, replacements={"data_root": data_folder},
    )
    valid_data = valid_data.filtered_sorted(sort_key="duration")

    test_real_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["csv_test_real"],
        replacements={"data_root": data_folder},
    )
    test_real_data = test_real_data.filtered_sorted(sort_key="duration")

    test_synth_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["csv_test_synth"],
        replacements={"data_root": data_folder},
    )
    test_synth_data = test_synth_data.filtered_sorted(sort_key="duration")

    all_real_data = sb.dataio.dataset.DynamicItemDataset.from_csv(
        csv_path=hparams["csv_all_real"],
        replacements={"data_root": data_folder},
    )
    all_real_data = all_real_data.filtered_sorted(sort_key="duration")

    datasets = [
        train_data,
        valid_data,
        test_real_data,
        test_synth_data,
        all_real_data,
    ]

    tokenizer = hparams["tokenizer"]

    # 2. Define audio pipeline:
    @sb.utils.data_pipeline.takes("wav")
    @sb.utils.data_pipeline.provides("sig")
    def audio_pipeline(wav):
        """ see original recipe
        """
        sig = sb.dataio.dataio.read_audio(wav)
        return sig

    sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)

    # 3. Define text pipeline:
    @sb.utils.data_pipeline.takes("semantics")
    @sb.utils.data_pipeline.provides(
        "semantics", "token_list", "tokens_bos", "tokens_eos", "tokens"
    )
    def text_pipeline(semantics):
        """ see original recipe
        """
        yield semantics
        tokens_list = tokenizer.encode_as_ids(semantics)
        yield tokens_list
        tokens_bos = torch.LongTensor([hparams["bos_index"]] + (tokens_list))
        yield tokens_bos
        tokens_eos = torch.LongTensor(tokens_list + [hparams["eos_index"]])
        yield tokens_eos
        tokens = torch.LongTensor(tokens_list)
        yield tokens

    sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline)

    # 4. Set output:
    sb.dataio.dataset.set_output_keys(
        datasets,
        [
            "id",
            "sig",
            "transcript",
            "semantics",
            "tokens_bos",
            "tokens_eos",
            "tokens",
        ],
    )
    return (
        train_data,
        valid_data,
        test_real_data,
        test_synth_data,
        all_real_data,
        tokenizer,
    )


if __name__ == "__main__":

    # Load hyperparameters file with command-line overrides
    hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
    with open(hparams_file) as fin:
        hparams = load_hyperpyyaml(fin, overrides)

    show_results_every = 100  # plots results every N iterations

    # If --distributed_launch then
    # create ddp_group with the right communication protocol
    sb.utils.distributed.ddp_init_group(run_opts)

    # Create experiment directory
    sb.create_experiment_directory(
        experiment_directory=hparams["output_folder"],
        hyperparams_to_save=hparams_file,
        overrides=overrides,
    )

    # Dataset prep (parsing TAS)
    """
    from prepare import prepare_TAS  # noqa

    # multi-gpu (ddp) save data preparation
    run_on_main(
        prepare_TAS,
        kwargs={
            "data_folder": hparams["data_folder"],
            "save_folder": hparams["output_folder"],
            "train_splits": hparams["train_splits"],
            "type": "decoupled",
            "skip_prep": hparams["skip_prep"],
        },
    )
    """

    # here we create the datasets objects as well as tokenization and encoding
    (
        train_set,
        valid_set,
        test_real_set,
        test_synth_set,
        all_real_set,
        tokenizer,
    ) = data_io_prepare(hparams)

    # We download and pretrain the tokenizer
    run_on_main(hparams["pretrainer"].collect_files)
    hparams["pretrainer"].load_collected(device=run_opts["device"])

    # Brain class initialization
    slu_brain = SLU(
        modules=hparams["modules"],
        opt_class=hparams["opt_class"],
        hparams=hparams,
        run_opts=run_opts,
        checkpointer=hparams["checkpointer"],
    )

    # adding objects to trainer:
    slu_brain.tokenizer = tokenizer

    # Training
    slu_brain.fit(
        slu_brain.hparams.epoch_counter,
        train_set,
        valid_set,
        train_loader_kwargs=hparams["dataloader_opts"],
        valid_loader_kwargs=hparams["dataloader_opts"],
    )

    # Test (ALL real data)
    if slu_brain.hparams.test_on_all_real:
        slu_brain.hparams.wer_file = (
            hparams["output_folder"] + "/wer_all_real.txt"
        )
        slu_brain.evaluate(
            all_real_set,
            test_loader_kwargs=hparams["dataloader_opts"],
            min_key="SER",
        )

    # Test (real data)
    slu_brain.hparams.wer_file = hparams["output_folder"] + "/wer_test_real.txt"
    slu_brain.evaluate(
        test_real_set,
        test_loader_kwargs=hparams["dataloader_opts"],
        min_key="SER",
    )

    # Test (synth data)
    slu_brain.hparams.wer_file = (
        hparams["output_folder"] + "/wer_test_synth.txt"
    )
    slu_brain.evaluate(
        test_synth_set,
        test_loader_kwargs=hparams["dataloader_opts"],
        min_key="SER",
    )

tests/templates/fetching_ddp_dynbatch_finetuning/multisource_mini_recipe.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""Recipe for debugging/testing, based upon SpeechBrain's minilibrispeech ASR template.

Does the following feature set work out together on some environment?
    DDP; dynamic batching; fine-tuning; mixed pretrainer fetching & testing using pretrained interface

Authors:
    * Andreas Nautsch 2023
"""
import logging
import speechbrain as sb
from copy import deepcopy
from speechbrain.pretrained import EncoderDecoderASR
from speechbrain.pretrained.fetching import FetchFrom, FetchSource


logger = logging.getLogger(__name__)


def run_hf_repo_single_node(run_options):
    """Test function for basic HF model use.

    Arguments
    ---------

    run_options: dict
        A set of options to change the runtime environment.
    """
    # Test on both nodes w/ pretrained HF model; ensure first it's fetched
    repo = FetchSource(
        FetchFrom.HUGGING_FACE,  # needs to be here only bc local path "speechbrain/asr-crdnn-rnnlm-librispeech" exists
        "speechbrain/asr-crdnn-rnnlm-librispeech",
    )
    EncoderDecoderASR.from_hparams(
        source=repo,  # to default save_dir: pretrained_models/...
        download_only=True,
    )

    # Instantiate pretrained HF model
    pretrained_hf_asr = EncoderDecoderASR.from_hparams(
        source=repo, run_opts=deepcopy(run_options),
    )

    # From HF model card
    pred = pretrained_hf_asr.transcribe_file(
        FetchSource(
            FetchFrom.HUGGING_FACE,  # bc HF-alike local path structure exists
            "speechbrain/asr-crdnn-rnnlm-librispeech/example.wav",
        )
    )

    print(f"prediction: {pred}")


def run_hf_other_repo_single_node(run_options):
    """Test function for basic HF model use.

    Arguments
    ---------

    run_options: dict
        A set of options to change the runtime environment.
    """
    # Test on both nodes w/ pretrained HF model; ensure first it's fetched
    repo = "speechbrain/asr-crdnn-commonvoice-fr"
    EncoderDecoderASR.from_hparams(
        source=repo,
        savedir="pretrained_models/asr-crdnn-commonvoice-fr",
        download_only=True,
    )

    # Instantiate pretrained HF model
    pretrained_hf_asr = EncoderDecoderASR.from_hparams(
        source=repo,
        savedir="pretrained_models/asr-crdnn-commonvoice-fr",
        run_opts=deepcopy(run_options),
    )

    # From HF model card
    pred = pretrained_hf_asr.transcribe_file(
        "speechbrain/asr-crdnn-commonvoice-fr/example-fr.wav"
    )

    print(f"prediction: {pred}")


if __name__ == "__main__":
    # hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])
    _, run_opts, _ = sb.parse_arguments(
        ["finetune.yaml", "--device", "cpu"]
    )  # just some yaml, so we get device into run_opts
    run_hf_repo_single_node(run_opts)
    run_hf_other_repo_single_node(run_opts)

tests/templates/fetching_ddp_dynbatch_finetuning/single_node_pretrained.py
TzuchengChang/NASS
import torch
import torch.nn


def test_SincConv(device):
    from speechbrain.nnet.CNN import SincConv

    input = torch.rand([4, 16000], device=device)
    convolve = SincConv(
        input_shape=input.shape, out_channels=8, kernel_size=65, padding="same"
    ).to(device)
    output = convolve(input)
    assert output.shape[-1] == 8

    assert torch.jit.trace(convolve, input)

    # Multichannel case
    input = torch.rand([10, 16000, 8], device=device)
    convolve = SincConv(
        input_shape=input.shape, out_channels=16, kernel_size=11, padding="same"
    ).to(device)
    output = convolve(input)
    assert output.shape[-1] == 16

    assert torch.jit.trace(convolve, input)


def test_Conv1d(device):

    from speechbrain.nnet.CNN import Conv1d

    input = (
        torch.tensor([-1, -1, -1, -1], device=device)
        .unsqueeze(0)
        .unsqueeze(2)
        .float()
    )
    convolve = Conv1d(
        out_channels=1, kernel_size=1, input_shape=input.shape, padding="same"
    ).to(device)
    output = convolve(input)
    assert input.shape == output.shape

    convolve.conv.weight = torch.nn.Parameter(
        torch.tensor([-1], device=device).float().unsqueeze(0).unsqueeze(1)
    )
    convolve.conv.bias = torch.nn.Parameter(
        torch.tensor([0], device=device).float()
    )
    output = convolve(input)
    assert torch.all(torch.eq(torch.ones(input.shape, device=device), output))

    assert torch.jit.trace(convolve, input)


def test_Conv2d(device):

    from speechbrain.nnet.CNN import Conv2d

    input = torch.rand([4, 11, 32, 1], device=device)
    convolve = Conv2d(
        out_channels=1,
        input_shape=input.shape,
        kernel_size=(1, 1),
        padding="same",
    ).to(device)
    output = convolve(input)
    assert output.shape[-1] == 1

    convolve.conv.weight = torch.nn.Parameter(
        torch.zeros(convolve.conv.weight.shape, device=device)
    )
    convolve.conv.bias = torch.nn.Parameter(
        torch.tensor([0], device=device).float()
    )
    output = convolve(input)
    assert torch.all(torch.eq(torch.zeros(input.shape, device=device), output))

    convolve.conv.weight = torch.nn.Parameter(
        torch.ones(convolve.conv.weight.shape, device=device)
    )
    convolve.conv.bias = torch.nn.Parameter(
        torch.tensor([0], device=device).float()
    )
    output = convolve(input)
    assert torch.all(torch.eq(input, output))

    assert torch.jit.trace(convolve, input)


def test_Leaf(device):
    from speechbrain.lobes.features import Leaf

    input = torch.rand([4, 16000], device=device)
    convolve = Leaf(
        input_shape=input.shape,
        window_len=25.0,
        window_stride=10.0,
        out_channels=8,
    ).to(device)
    output = convolve(input)
    assert output.shape[-1] == 8
    assert torch.jit.trace(convolve, input)

tests/unittests/test_CNN.py
TzuchengChang/NASS
import torch
import torch.nn
from collections import OrderedDict


def test_RNN(device):

    from speechbrain.nnet.RNN import RNN, GRU, LSTM, LiGRU, QuasiRNN, RNNCell

    # Check RNN
    inputs = torch.randn(4, 2, 7, device=device)
    net = RNN(
        hidden_size=5,
        input_shape=inputs.shape,
        num_layers=2,
        bidirectional=False,
    ).to(device)
    output, hn = net(inputs)
    output_l = []
    hn_t = None
    for t in range(inputs.shape[1]):
        out_t, hn_t = net(inputs[:, t, :].unsqueeze(1), hn_t)
        output_l.append(out_t.squeeze(1))

    out_steps = torch.stack(output_l, dim=1)
    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "GRU output mismatch"
    assert torch.all(
        torch.lt(torch.add(hn_t, -hn), 1e-3)
    ), "GRU hidden states mismatch"
    assert torch.jit.trace(net, inputs)

    # Check GRU
    inputs = torch.randn(4, 2, 7, device=device)
    net = GRU(
        hidden_size=5,
        input_shape=inputs.shape,
        num_layers=2,
        bidirectional=False,
    ).to(device)
    output, hn = net(inputs)
    output_l = []
    hn_t = None
    for t in range(inputs.shape[1]):
        out_t, hn_t = net(inputs[:, t, :].unsqueeze(1), hn_t)
        output_l.append(out_t.squeeze(1))

    out_steps = torch.stack(output_l, dim=1)
    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "GRU output mismatch"
    assert torch.all(
        torch.lt(torch.add(hn_t, -hn), 1e-3)
    ), "GRU hidden states mismatch"
    assert torch.jit.trace(net, inputs)

    # Check LSTM
    inputs = torch.randn(4, 2, 7, device=device)
    net = LSTM(
        hidden_size=5,
        input_shape=inputs.shape,
        num_layers=2,
        bidirectional=False,
    ).to(device)
    output, hn = net(inputs)
    output_l = []
    hn_t = None
    for t in range(inputs.shape[1]):
        out_t, hn_t = net(inputs[:, t, :].unsqueeze(1), hn_t)
        output_l.append(out_t.squeeze(1))

    out_steps = torch.stack(output_l, dim=1)
    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "LSTM output mismatch"
    assert torch.all(torch.lt(torch.add(hn_t[0], -hn[0]), 1e-3)) and torch.all(
        torch.lt(torch.add(hn_t[1], -hn[1]), 1e-3)
    ), "LSTM hidden states mismatch"
    assert torch.jit.trace(net, inputs)

    # Check LiGRU
    inputs = torch.randn(1, 2, 2, device=device)
    net = LiGRU(
        hidden_size=5,
        input_shape=inputs.shape,
        num_layers=2,
        bidirectional=False,
        normalization="layernorm",
    ).to(device)

    output, hn = net(inputs)
    output_l = []
    hn_t = None
    for t in range(inputs.shape[1]):
        out_t, hn_t = net(inputs[:, t, :].unsqueeze(1), hn_t)
        output_l.append(out_t.squeeze(1))

    out_steps = torch.stack(output_l, dim=1)

    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "LiGRU output mismatch"
    assert torch.all(torch.lt(torch.add(hn_t[0], -hn[0]), 1e-3)) and torch.all(
        torch.lt(torch.add(hn_t[1], -hn[1]), 1e-3)
    ), "LiGRU hidden states mismatch"

    # Check QuasiRNN
    inputs = torch.randn(1, 2, 2, device=device)
    net = QuasiRNN(
        hidden_size=5,
        input_shape=inputs.shape,
        num_layers=2,
        bidirectional=False,
    ).to(device)

    output, hn = net(inputs)
    output_l = []
    hn_t = None
    for t in range(inputs.shape[1]):
        out_t, hn_t = net(inputs[:, t, :].unsqueeze(1), hn_t)
        output_l.append(out_t.squeeze(1))

    out_steps = torch.stack(output_l, dim=1)

    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "QuasiRNN output mismatch"
    assert torch.all(
        torch.lt(torch.add(hn_t[0], -hn[0][1]), 1e-3)
    ) and torch.all(
        torch.lt(torch.add(hn_t[1], -hn[1][1]), 1e-3)
    ), "QuasiRNN hidden states mismatch"
    assert torch.jit.trace(net, inputs)

    # Check RNNCell
    inputs = torch.randn(4, 2, 7, device=device)
    net = RNNCell(hidden_size=5, input_size=7, num_layers=2, dropout=0.0).to(
        device
    )
    hn_t = None
    output_lst = []
    for t in range(inputs.shape[1]):
        output, hn_t = net(inputs[:, t], hn_t)
        output_lst.append(output)

    out_steps = torch.stack(output_lst, dim=1)
    rnn = torch.nn.RNN(
        input_size=7, hidden_size=5, num_layers=2, batch_first=True
    ).to(device)

    # rename the state_dict
    state = net.state_dict()
    new_state = []
    for name, tensor in state.items():
        index, weight_id = name[len("rnn_cells.")], name[len("rnn_cells.0.") :]
        new_state.append((f"{weight_id}_l{index}", tensor))
    new_state = OrderedDict(new_state)
    rnn.load_state_dict(new_state)
    output, hn = rnn(inputs)

    assert torch.all(
        torch.lt(torch.add(out_steps, -output), 1e-3)
    ), "RNNCell output mismatch"
    assert torch.all(torch.lt(torch.add(hn_t[0], -hn[0]), 1e-3)) and torch.all(
        torch.lt(torch.add(hn_t[1], -hn[1]), 1e-3)
    ), "RNNCell hidden states mismatch"
    assert torch.jit.trace(rnn, inputs)

tests/unittests/test_RNN.py
TzuchengChang/NASS
import torch
import torch.nn


def test_softmax(device):

    from speechbrain.nnet.activations import Softmax

    inputs = torch.tensor([1, 2, 3], device=device).float()
    act = Softmax(apply_log=False)
    outputs = act(inputs)
    assert torch.argmax(outputs) == 2

    assert torch.jit.trace(act, inputs)

tests/unittests/test_activations.py
TzuchengChang/NASS
import torch


def test_rel_pos_MHA(device):

    from speechbrain.nnet.attention import RelPosMHAXL

    bsz = 2
    emb_dim = 4
    k_len = [12, 10]
    q_len = [10, 12]
    bias = [True, False]
    head_dim = [4, None]

    for kl in k_len:
        for ql in q_len:
            for b in bias:
                for h in head_dim:
                    relpos = RelPosMHAXL(
                        emb_dim, num_heads=2, vbias=b, vdim=h
                    ).to(device)
                    q = torch.rand((bsz, ql, emb_dim), device=device)
                    k = torch.rand((bsz, kl, emb_dim), device=device)
                    pos_embs = torch.rand(
                        (1, 2 * kl - 1, emb_dim), device=device
                    )
                    relpos(q, k, k, pos_embs=pos_embs)

tests/unittests/test_attention.py
TzuchengChang/NASS
import os
import torch
from speechbrain.dataio.dataio import write_audio


def test_add_noise(tmpdir, device):
    from speechbrain.processing.speech_augmentation import AddNoise

    # Test concatenation of batches
    wav_a = torch.sin(torch.arange(8000.0, device=device)).unsqueeze(0)
    a_len = torch.ones(1, device=device)
    wav_b = (
        torch.cos(torch.arange(10000.0, device=device))
        .unsqueeze(0)
        .repeat(2, 1)
    )
    b_len = torch.ones(2, device=device)
    concat, lens = AddNoise._concat_batch(wav_a, a_len, wav_b, b_len)
    assert concat.shape == (3, 10000)
    assert lens.allclose(torch.Tensor([0.8, 1, 1]).to(device))
    concat, lens = AddNoise._concat_batch(wav_b, b_len, wav_a, a_len)
    assert concat.shape == (3, 10000)
    expected = torch.Tensor([1, 1, 0.8]).to(device)
    assert lens.allclose(expected)

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)
    test_noise = torch.cos(torch.arange(16000.0, device=device)).unsqueeze(0)
    wav_lens = torch.ones(1, device=device)

    # Put noise waveform into temporary file
    noisefile = os.path.join(tmpdir, "noise.wav")
    write_audio(noisefile, test_noise.transpose(0, 1).cpu(), 16000)

    csv = os.path.join(tmpdir, "noise.csv")
    with open(csv, "w") as w:
        w.write("ID, duration, wav, wav_format, wav_opts\n")
        w.write(f"1, 1.0, {noisefile}, wav,\n")

    # Edge cases
    no_noise = AddNoise(mix_prob=0.0).to(device)
    assert no_noise(test_waveform, wav_lens).allclose(test_waveform)
    no_noise = AddNoise(snr_low=1000, snr_high=1000)
    assert no_noise(test_waveform, wav_lens).allclose(test_waveform)
    all_noise = AddNoise(csv_file=csv, snr_low=-1000, snr_high=-1000)
    assert all_noise(test_waveform, wav_lens).allclose(test_noise, atol=1e-4)

    # Basic 0dB case
    add_noise = AddNoise(csv_file=csv).to(device)
    expected = (test_waveform + test_noise) / 2
    assert add_noise(test_waveform, wav_lens).allclose(expected, atol=1e-4)


def test_add_reverb(tmpdir, device):
    from speechbrain.processing.speech_augmentation import AddReverb

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)
    impulse_response = torch.zeros(1, 8000, device=device)
    impulse_response[0, 0] = 1.0
    wav_lens = torch.ones(1, device=device)

    # Put ir waveform into temporary file
    ir1 = os.path.join(tmpdir, "ir1.wav")
    ir2 = os.path.join(tmpdir, "ir2.wav")
    ir3 = os.path.join(tmpdir, "ir3.wav")
    write_audio(ir1, impulse_response.cpu().transpose(0, 1), 16000)

    impulse_response[0, 0] = 0.0
    impulse_response[0, 10] = 0.5
    write_audio(ir2, impulse_response.cpu().transpose(0, 1), 16000)

    # Check a very simple non-impulse-response case:
    impulse_response[0, 10] = 0.6
    impulse_response[0, 11] = 0.4
    # sf.write(ir3, impulse_response.squeeze(0).numpy(), 16000)
    write_audio(ir3, impulse_response.cpu().transpose(0, 1), 16000)
    ir3_result = test_waveform * 0.6 + test_waveform.roll(1, -1) * 0.4

    # write ir csv file
    csv = os.path.join(tmpdir, "ir.csv")
    with open(csv, "w") as w:
        w.write("ID, duration, wav, wav_format, wav_opts\n")
        w.write(f"1, 0.5, {ir1}, wav,\n")
        w.write(f"2, 0.5, {ir2}, wav,\n")
        w.write(f"3, 0.5, {ir3}, wav,\n")

    # Edge case
    no_reverb = AddReverb(csv, reverb_prob=0.0).to(device)
    assert no_reverb(test_waveform, wav_lens).allclose(test_waveform)

    # Normal cases
    add_reverb = AddReverb(csv, sorting="original")
    reverbed = add_reverb(test_waveform, wav_lens)[:, 0:1000]
    assert reverbed.allclose(test_waveform[:, 0:1000], atol=1e-1)
    reverbed = add_reverb(test_waveform, wav_lens)[:, 0:1000]
    assert reverbed.allclose(test_waveform[:, 0:1000], atol=1e-1)
    reverbed = add_reverb(test_waveform, wav_lens)[:, 0:1000]
    assert reverbed.allclose(ir3_result[:, 0:1000], atol=2e-1)


def test_speed_perturb(device):
    from speechbrain.processing.speech_augmentation import SpeedPerturb

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)

    # Edge cases
    no_perturb = SpeedPerturb(16000, perturb_prob=0.0).to(device)
    assert no_perturb(test_waveform).allclose(test_waveform)
    no_perturb = SpeedPerturb(16000, speeds=[100]).to(device)
    assert no_perturb(test_waveform).allclose(test_waveform)

    # Half speed
    half_speed = SpeedPerturb(16000, speeds=[50]).to(device)
    assert half_speed(test_waveform).allclose(test_waveform[:, ::2], atol=3e-1)


def test_babble(device):
    from speechbrain.processing.speech_augmentation import AddBabble

    test_waveform = torch.stack(
        (
            torch.sin(torch.arange(16000.0, device=device)),
            torch.cos(torch.arange(16000.0, device=device)),
        )
    )
    lengths = torch.ones(2, device=device)

    # Edge cases
    no_babble = AddBabble(mix_prob=0.0).to(device)
    assert no_babble(test_waveform, lengths).allclose(test_waveform)
    no_babble = AddBabble(speaker_count=1, snr_low=1000, snr_high=1000)
    assert no_babble(test_waveform, lengths).allclose(test_waveform)

    # One babbler just averages the two speakers
    babble = AddBabble(speaker_count=1).to(device)
    expected = (test_waveform + test_waveform.roll(1, 0)) / 2
    assert babble(test_waveform, lengths).allclose(expected, atol=1e-4)


def test_drop_freq(device):
    from speechbrain.processing.speech_augmentation import DropFreq

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)

    # Edge cases
    no_drop = DropFreq(drop_prob=0.0).to(device)
    assert no_drop(test_waveform).allclose(test_waveform)
    no_drop = DropFreq(drop_count_low=0, drop_count_high=0)
    assert no_drop(test_waveform).allclose(test_waveform)

    # Check case where frequency range *does not* include signal frequency
    drop_diff_freq = DropFreq(drop_freq_low=0.5, drop_freq_high=0.9)
    assert drop_diff_freq(test_waveform).allclose(test_waveform, atol=1e-1)

    # Check case where frequency range *does* include signal frequency
    drop_same_freq = DropFreq(drop_freq_low=0.28, drop_freq_high=0.28)
    assert drop_same_freq(test_waveform).allclose(
        torch.zeros(1, 16000, device=device), atol=4e-1
    )


def test_drop_chunk(device):
    from speechbrain.processing.speech_augmentation import DropChunk

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)
    lengths = torch.ones(1, device=device)

    # Edge cases
    no_drop = DropChunk(drop_prob=0.0).to(device)
    assert no_drop(test_waveform, lengths).allclose(test_waveform)
    no_drop = DropChunk(drop_length_low=0, drop_length_high=0).to(device)
    assert no_drop(test_waveform, lengths).allclose(test_waveform)
    no_drop = DropChunk(drop_count_low=0, drop_count_high=0).to(device)
    assert no_drop(test_waveform, lengths).allclose(test_waveform)
    no_drop = DropChunk(drop_start=0, drop_end=0).to(device)
    assert no_drop(test_waveform, lengths).allclose(test_waveform)

    # Specify all parameters to ensure it is deterministic
    dropper = DropChunk(
        drop_length_low=100,
        drop_length_high=100,
        drop_count_low=1,
        drop_count_high=1,
        drop_start=100,
        drop_end=200,
        noise_factor=0.0,
    ).to(device)
    expected_waveform = test_waveform.clone()
    expected_waveform[:, 100:200] = 0.0

    assert dropper(test_waveform, lengths).allclose(expected_waveform)

    # Make sure amplitude is similar before and after
    dropper = DropChunk(noise_factor=1.0).to(device)
    drop_amplitude = dropper(test_waveform, lengths).abs().mean()
    orig_amplitude = test_waveform.abs().mean()
    assert drop_amplitude.allclose(orig_amplitude, atol=1e-2)


def test_clip(device):
    from speechbrain.processing.speech_augmentation import DoClip

    test_waveform = torch.sin(torch.arange(16000.0, device=device)).unsqueeze(0)

    # Edge cases
    no_clip = DoClip(clip_prob=0.0).to(device)
    assert no_clip(test_waveform).allclose(test_waveform)
    no_clip = DoClip(clip_low=1, clip_high=1).to(device)
    assert no_clip(test_waveform).allclose(test_waveform)

    # Sort of a reimplementation of clipping, but its one function call.
    expected = test_waveform.clamp(min=-0.5, max=0.5)
    half_clip = DoClip(clip_low=0.5, clip_high=0.5).to(device)
    assert half_clip(test_waveform).allclose(expected)

tests/unittests/test_augment.py
TzuchengChang/NASS
import pytest
import torch
import numpy as np


def test_batch_pad_right_to(device):
    from speechbrain.utils.data_utils import batch_pad_right
    import random

    n_channels = 40
    batch_lens = [1, 5]

    for b in batch_lens:
        rand_lens = [random.randint(10, 53) for x in range(b)]
        tensors = [
            torch.ones((rand_lens[x], n_channels), device=device)
            for x in range(b)
        ]
        batched, lens = batch_pad_right(tensors)
        assert batched.shape[0] == b
        np.testing.assert_almost_equal(
            lens, [x / max(rand_lens) for x in rand_lens], decimal=3
        )

    for b in batch_lens:
        rand_lens = [random.randint(10, 53) for x in range(b)]
        tensors = [torch.ones(rand_lens[x], device=device) for x in range(b)]
        batched, lens = batch_pad_right(tensors)
        assert batched.shape[0] == b
        np.testing.assert_almost_equal(
            lens, [x / max(rand_lens) for x in rand_lens], decimal=3
        )


def test_paddedbatch(device):
    from speechbrain.dataio.batch import PaddedBatch

    batch = PaddedBatch(
        [
            {
                "id": "ex1",
                "foo": torch.Tensor([1.0]).to(device),
                "bar": torch.Tensor([1.0, 2.0, 3.0]).to(device),
            },
            {
                "id": "ex2",
                "foo": torch.Tensor([2.0, 1.0]).to(device),
                "bar": torch.Tensor([2.0]).to(device),
            },
        ]
    )
    batch.to(dtype=torch.half)
    assert batch.foo.data.dtype == torch.half
    assert batch["foo"][1].dtype == torch.half
    assert batch.bar.lengths.dtype == torch.half
    assert batch.foo.data.shape == torch.Size([2, 2])
    assert batch.bar.data.shape == torch.Size([2, 3])
    ids, foos, bars = batch
    assert ids == ["ex1", "ex2"]


@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
def test_pin_memory():
    from speechbrain.dataio.batch import PaddedBatch

    batch = PaddedBatch(
        [
            {
                "id": "ex1",
                "foo": torch.Tensor([1.0]),
                "bar": torch.Tensor([1.0, 2.0, 3.0]),
            },
            {
                "id": "ex2",
                "foo": torch.Tensor([2.0, 1.0]),
                "bar": torch.Tensor([2.0]),
            },
        ]
    )
    batch.pin_memory()
    assert batch.foo.data.is_pinned()

tests/unittests/test_batching.py
TzuchengChang/NASS
import pytest


def test_categorical_encoder(device):
    from speechbrain.dataio.encoder import CategoricalEncoder

    encoder = CategoricalEncoder()
    encoder.expect_len(4)
    encoder.update_from_iterable("abcd")
    integers = encoder.encode_sequence("dcba")
    assert all(isinstance(i, int) for i in integers)
    assert encoder.is_continuous()
    with pytest.raises(KeyError):
        encoder.add_label("a")
    # Does NOT raise:
    encoder.ensure_label("a")
    with pytest.raises(KeyError):
        encoder.insert_label("a", -3)
    encoder.enforce_label("a", -3)
    assert encoder.encode_label("a") == -3
    assert not encoder.is_continuous()

    # Decoding:
    import torch

    encoder = CategoricalEncoder()
    encoder.expect_len(4)
    encoder.update_from_iterable("abcd")
    result = encoder.decode_torch(
        torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3]], device=device)
    )
    assert result == [["a", "a"], ["b", "b"], ["c", "c"], ["d", "d"]]
    result = encoder.decode_ndim([[0, 0], [1, 1], [2, 2], [3, 3]])
    assert result == [["a", "a"], ["b", "b"], ["c", "c"], ["d", "d"]]
    result = encoder.decode_ndim(torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3]]))
    assert result == [["a", "a"], ["b", "b"], ["c", "c"], ["d", "d"]]
    result = encoder.decode_ndim([[[[[0, 0], [1, 1], [2, 2], [3, 3]]]]])
    assert result == [[[[["a", "a"], ["b", "b"], ["c", "c"], ["d", "d"]]]]]
    result = encoder.decode_torch(
        torch.tensor([[[[[0, 0], [1, 1], [2, 2], [3, 3]]]]], device=device)
    )
    assert result == [[[[["a", "a"], ["b", "b"], ["c", "c"], ["d", "d"]]]]]
    result = encoder.decode_ndim([[0, 0], [1], [2, 2, 2], []])
    assert result == [["a", "a"], ["b"], ["c", "c", "c"], []]

    encoder = CategoricalEncoder()
    encoder.expect_len(3)
    encoder.limited_labelset_from_iterable("aabbbcccd", n_most_common=3)
    encoder.encode_sequence("abc")
    with pytest.raises(KeyError):
        encoder.encode_label("d")
    encoder = CategoricalEncoder()
    encoder.expect_len(2)
    encoder.limited_labelset_from_iterable("aabbbcccd", min_count=3)
    encoder.encode_sequence("cbcb")
    with pytest.raises(KeyError):
        encoder.encode_label("a")
    with pytest.raises(KeyError):
        encoder.encode_label("d")
    encoder = CategoricalEncoder()
    encoder.expect_len(2)
    encoder.limited_labelset_from_iterable(
        "aabbbcccd", n_most_common=3, min_count=3
    )
    encoder.encode_sequence("cbcb")
    with pytest.raises(KeyError):
        encoder.encode_label("a")
    with pytest.raises(KeyError):
        encoder.encode_label("d")

    encoder = CategoricalEncoder(unk_label="<unk>")
    encoder.expect_len(4)
    encoder.update_from_iterable("abc")
    assert encoder.encode_label("a") == 1
    assert encoder.encode_label("d") == 0
    assert encoder.decode_ndim(encoder.encode_label("d")) == "<unk>"


def test_categorical_encoder_saving(tmpdir):
    from speechbrain.dataio.encoder import CategoricalEncoder

    encoder = CategoricalEncoder(starting_index=3)
    encoding_file = tmpdir / "char_encoding.txt"
    # First time this runs, the encoding is created:
    if not encoder.load_if_possible(encoding_file):
        encoder.update_from_iterable("abcd")
        encoder.save(encoding_file)
    else:
        assert False  # We should not get here!
    # Now, imagine a recovery:
    encoder = CategoricalEncoder()
    encoder.expect_len(4)
    # The second time, the encoding is just loaded from file:
    if not encoder.load_if_possible(encoding_file):
        assert False  # We should not get here!
    integers = encoder.encode_sequence("dcba")
    assert all(isinstance(i, int) for i in integers)
    assert encoder.starting_index == 3  # This is also loaded

    # Also possible to encode tuples and load
    encoder = CategoricalEncoder()
    encoding_file = tmpdir / "tuple_encoding.txt"
    encoder.add_label((1, 2, 3))
    encoder.insert_label((1, 2), index=-1)
    encoder.save(encoding_file)
    # Reload
    encoder = CategoricalEncoder()
    encoder.expect_len(2)
    assert encoder.load_if_possible(encoding_file)
    assert encoder.encode_label((1, 2)) == -1

    # Load unk:
    encoder = CategoricalEncoder(unk_label="UNKNOWN")
    encoding_file = tmpdir / "unk_encoding.txt"
    encoder.update_from_iterable("abc")
    encoder.save(encoding_file)
    encoder = CategoricalEncoder()
    encoder.expect_len(4)
    assert encoder.load_if_possible(encoding_file)
    assert encoder.encode_label("a") == 1
    assert encoder.decode_ndim(encoder.encode_label("d")) == "UNKNOWN"
    # Even if set differently:
    encoder = CategoricalEncoder()
    encoder.add_unk()
    encoder.expect_len(4)
    assert encoder.load_if_possible(encoding_file)
    assert encoder.encode_label("a") == 1
    assert encoder.decode_ndim(encoder.encode_label("d")) == "UNKNOWN"


def test_categorical_encoder_from_dataset():
    from speechbrain.dataio.encoder import CategoricalEncoder
    from speechbrain.dataio.dataset import DynamicItemDataset

    encoder = CategoricalEncoder()
    data = {
        "utt1": {"foo": -1, "bar": 0, "text": "hello world"},
        "utt2": {"foo": 1, "bar": 2, "text": "how are you world"},
        "utt3": {"foo": 3, "bar": 4, "text": "where are you world"},
        "utt4": {"foo": 5, "bar": 6, "text": "hello nation"},
    }
    dynamic_items = [
        {"func": lambda x: x.split(), "takes": ["text"], "provides": "words"},
        {
            "func": encoder.encode_sequence,
            "takes": ["words"],
            "provides": "words_t",
        },
    ]
    output_keys = ["words_t"]
    dataset = DynamicItemDataset(data, dynamic_items, output_keys)
    encoder.update_from_didataset(dataset, "words", sequence_input=True)
    encoder.expect_len(7)
    assert dataset[0]["words_t"] == [0, 1]
    assert encoder.decode_ndim(dataset[0]["words_t"]) == ["hello", "world"]


def test_categorical_encoder_length_check():
    from speechbrain.dataio.encoder import CategoricalEncoder

    encoder = CategoricalEncoder()
    encoder.update_from_iterable("abcd")

    encoder.expect_len(3)
    with pytest.raises(RuntimeError):
        encoder.encode_label("a")

    encoder.ignore_len()
    encoder.encode_label("a")

    encoder.expect_len(4)
    encoder.encode_label("a")


def test_text_encoder(tmpdir):
    from speechbrain.dataio.encoder import TextEncoder

    encoder = TextEncoder()
    encoding_file = tmpdir / "text_encoding.txt"
    encoder.add_bos_eos()
    encoder.update_from_iterable(
        [["hello", "world"], ["how", "are", "you", "world"]],
        sequence_input=True,
    )
    encoder.expect_len(7)
    encoded = encoder.encode_sequence(
        encoder.prepend_bos_label(["are", "you", "world"])
    )
    assert encoded[0] == 0
    encoded = encoder.append_eos_index(
        encoder.encode_sequence(["are", "you", "world"])
    )
    assert encoded[-1] == 1  # By default uses just one sentence_boundary marker
    encoder.save(encoding_file)
    encoder = TextEncoder()
    encoder.expect_len(7)
    assert encoder.load_if_possible(encoding_file)
    encoded = encoder.encode_sequence(
        encoder.append_eos_label(["are", "you", "world"])
    )
    assert encoded[-1] == 1
    encoded = encoder.prepend_bos_index(
        encoder.encode_sequence(["are", "you", "world"])
    )
    assert encoded[0] == 0


def test_ctc_encoder(tmpdir):
    from speechbrain.dataio.encoder import CTCTextEncoder

    encoder = CTCTextEncoder()
    encoder.expect_len(9)  # "abcdef" + bos + eos + blank
    encoder.insert_bos_eos(
        bos_label="<s>", bos_index=0, eos_label="</s>", eos_index=1
    )
    encoder.insert_blank(blank_label="_", index=2)
    encoding_file = tmpdir / "ctc_encoding.txt"
    encoder.update_from_iterable(["abcd", "bcdef"], sequence_input=True)
    encoded = encoder.encode_sequence(encoder.prepend_bos_label(["a", "b"]))
    assert encoded[0] == 0
    encoder.save(encoding_file)
    encoder = CTCTextEncoder()
    encoder.expect_len(9)
    assert encoder.load_if_possible(encoding_file)
    assert (
        "".join(encoder.collapse_labels("_bb_aaa___bbbbb_b_eeee_____"))
        == "babbe"
    )
    assert "".join(encoder.collapse_labels("babe")) == "babe"
    assert (
        "".join(
            encoder.collapse_labels(
                "_bb_aaa___bbbbb_b_eeee_____", merge_repeats=False
            )
        )
        == "bbaaabbbbbbeeee"
    )
    assert encoder.decode_ndim(
        (
            encoder.collapse_indices_ndim(
                [
                    [0, 2, 4, 4, 2, 3, 3, 3, 2, 2, 2, 4, 2, 4, 2, 7, 2, 1],
                    [[0, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1]],
                ]
            )
        )
    ) == [
        ["<s>", "b", "a", "b", "b", "e", "</s>"],
        [["<s>", "a", "b", "c", "d", "c", "b", "a", "</s>"]],
    ]

tests/unittests/test_categorical_encoder.py
TzuchengChang/NASS
import pytest


def test_checkpointer(tmpdir, device):
    from speechbrain.utils.checkpoints import Checkpointer
    import torch

    class Recoverable(torch.nn.Module):
        def __init__(self, param):
            super().__init__()
            self.param = torch.nn.Parameter(torch.tensor([param]))

        def forward(self, x):
            return x * self.param

    recoverable = Recoverable(2.0)
    recoverables = {"recoverable": recoverable}
    recoverer = Checkpointer(tmpdir, recoverables)
    recoverable.param.data = torch.tensor([1.0], device=device)
    # Should not be possible since no checkpoint saved yet:
    assert not recoverer.recover_if_possible()
    result = recoverable(10.0)
    # Check that parameter has not been loaded from original value:
    assert recoverable.param.data == torch.tensor([1.0], device=device)

    ckpt = recoverer.save_checkpoint()
    # Check that the name recoverable has a save file:
    # NOTE: Here assuming .pt filename; if convention changes, change test
    assert (ckpt.path / "recoverable.ckpt").exists()
    # Check that saved checkpoint is found, and location correct:
    assert recoverer.list_checkpoints()[0] == ckpt
    assert recoverer.list_checkpoints()[0].path.parent == tmpdir
    recoverable.param.data = torch.tensor([2.0], device=device)
    recoverer.recover_if_possible()
    # Check that parameter has been loaded immediately:
    assert recoverable.param.data == torch.tensor([1.0], device=device)
    result = recoverable(10.0)
    # And result correct
    assert result == 10.0

    other = Recoverable(2.0)
    recoverer.add_recoverable("other", other)
    # Check that both objects are now found:
    assert recoverer.recoverables["recoverable"] == recoverable
    assert recoverer.recoverables["other"] == other
    new_ckpt = recoverer.save_checkpoint()
    # Check that now both recoverables have a save file:
    assert (new_ckpt.path / "recoverable.ckpt").exists()
    assert (new_ckpt.path / "other.ckpt").exists()
    assert new_ckpt in recoverer.list_checkpoints()
    recoverable.param.data = torch.tensor([2.0], device=device)
    other.param.data = torch.tensor([10.0], device=device)
    chosen_ckpt = recoverer.recover_if_possible()
    # Should choose newest by default:
    assert chosen_ckpt == new_ckpt
    # Check again that parameters have been loaded immediately:
    assert recoverable.param.data == torch.tensor([1.0], device=device)
    assert other.param.data == torch.tensor([2.0], device=device)
    other_result = other(10.0)
    # And again we should have the correct computations:
    assert other_result == 20.0

    # Recover from oldest, which does not have "other":
    # This also tests a custom sort
    # Raises by default:
    with pytest.raises(RuntimeError):
        chosen_ckpt = recoverer.recover_if_possible(
            importance_key=lambda x: -x.meta["unixtime"]
        )
    # However this operation may have loaded the first object
    # so let's set the values manually:
    recoverable.param.data = torch.tensor([2.0], device=device)
    other.param.data = torch.tensor([10.0], device=device)
    recoverer.allow_partial_load = True
    chosen_ckpt = recoverer.recover_if_possible(
        importance_key=lambda x: -x.meta["unixtime"]
    )
    # Should have chosen the original:
    assert chosen_ckpt == ckpt
    # And should recover recoverable:
    assert recoverable.param.data == torch.tensor([1.0], device=device)
    # But not other:
    other_result = other(10.0)
    assert other.param.data == torch.tensor([10.0], device=device)
    assert other_result == 100.0

    # Test saving names checkpoints with meta info, and custom filter
    epoch_ckpt = recoverer.save_checkpoint(name="ep1", meta={"loss": 2.0})
    assert "ep1" in epoch_ckpt.path.name
    other.param.data = torch.tensor([2.0], device=device)
    recoverer.save_checkpoint(meta={"loss": 3.0})
    chosen_ckpt = recoverer.recover_if_possible(
        ckpt_predicate=lambda ckpt: "loss" in ckpt.meta,
        importance_key=lambda ckpt: -ckpt.meta["loss"],
    )
    assert chosen_ckpt == epoch_ckpt
    assert other.param.data == torch.tensor([10.0], device=device)

    # Make sure checkpoints can't be name saved by the same name
    with pytest.raises(FileExistsError):
        recoverer.save_checkpoint(name="ep1")


def test_recovery_custom_io(tmpdir):
    from speechbrain.utils.checkpoints import register_checkpoint_hooks
    from speechbrain.utils.checkpoints import mark_as_saver
    from speechbrain.utils.checkpoints import mark_as_loader
    from speechbrain.utils.checkpoints import Checkpointer

    @register_checkpoint_hooks
    class CustomRecoverable:
        def __init__(self, param):
            self.param = int(param)

        @mark_as_saver
        def save(self, path):
            with open(path, "w") as fo:
                fo.write(str(self.param))

        @mark_as_loader
        def load(self, path, end_of_epoch, device):
            del end_of_epoch  # Unused
            del device
            with open(path) as fi:
                self.param = int(fi.read())

    custom_recoverable = CustomRecoverable(0)
    recoverer = Checkpointer(tmpdir, {"custom_recoverable": custom_recoverable})
    custom_recoverable.param = 1
    # First, make sure no checkpoints are found
    # (e.g. somehow tmpdir contaminated)
    ckpt = recoverer.recover_if_possible()
    assert ckpt is None
    ckpt = recoverer.save_checkpoint()
    custom_recoverable.param = 2
    loaded_ckpt = recoverer.recover_if_possible()
    # Make sure we got the same thing:
    assert ckpt == loaded_ckpt
    # With this custom recoverable, the load is instant:
    assert custom_recoverable.param == 1


def test_checkpoint_deletion(tmpdir, device):
    from speechbrain.utils.checkpoints import Checkpointer
    import torch

    class Recoverable(torch.nn.Module):
        def __init__(self, param):
            super().__init__()
            self.param = torch.nn.Parameter(
                torch.tensor([param], device=device)
            )

        def forward(self, x):
            return x * self.param

    recoverable = Recoverable(1.0)
    recoverables = {"recoverable": recoverable}
    recoverer = Checkpointer(tmpdir, recoverables)
    first_ckpt = recoverer.save_checkpoint()
    recoverer.delete_checkpoints()
    # Will not delete only checkpoint by default:
    assert first_ckpt in recoverer.list_checkpoints()
    second_ckpt = recoverer.save_checkpoint()
    recoverer.delete_checkpoints()
    # Oldest checkpoint is deleted by default:
    assert first_ckpt not in recoverer.list_checkpoints()
    # Other syntax also should work:
    recoverer.save_and_keep_only()
    assert second_ckpt not in recoverer.list_checkpoints()
    # Can delete all checkpoints:
    recoverer.delete_checkpoints(num_to_keep=0)
    assert not recoverer.list_checkpoints()

    # Now each should be kept:
    # Highest foo
    c1 = recoverer.save_checkpoint(meta={"foo": 2})
    # Latest CKPT after filtering
    c2 = recoverer.save_checkpoint(meta={"foo": 1})
    # Filtered out
    c3 = recoverer.save_checkpoint(meta={"epoch_ckpt": True})
    recoverer.delete_checkpoints(
        num_to_keep=1,
        max_keys=["foo"],
        importance_keys=[lambda c: c.meta["unixtime"]],
        ckpt_predicate=lambda c: "epoch_ckpt" not in c.meta,
    )
    assert all(c in recoverer.list_checkpoints() for c in [c1, c2, c3])
    # Reset:
    recoverer.delete_checkpoints(num_to_keep=0)
    assert not recoverer.list_checkpoints()

    # Test the keeping multiple checkpoints without predicate:
    # This should be deleted:
    c_to_delete = recoverer.save_checkpoint(meta={"foo": 2})
    # Highest foo
    c1 = recoverer.save_checkpoint(meta={"foo": 3})
    # Latest CKPT after filtering
    c2 = recoverer.save_checkpoint(meta={"foo": 1})
    recoverer.delete_checkpoints(
        num_to_keep=1,
        importance_keys=[lambda c: c.meta["unixtime"], lambda c: c.meta["foo"]],
    )
    assert all(c in recoverer.list_checkpoints() for c in [c1, c2])
    assert c_to_delete not in recoverer.list_checkpoints()


def test_multiple_ckpts_and_criteria(tmpdir):
    from speechbrain.utils.checkpoints import Checkpointer
    import torch

    class Recoverable(torch.nn.Module):
        def __init__(self, param):
            super().__init__()
            self.param = torch.nn.Parameter(torch.tensor([param]))

        def forward(self, x):
            return x * self.param

    recoverable = Recoverable(1.0)
    recoverables = {"recoverable": recoverable}
    recoverer = Checkpointer(tmpdir, recoverables)

    # Here testing multiple checkpoints with equal meta criteria
    recoverer.save_and_keep_only(
        meta={"error": 5}, min_keys=["error"], keep_recent=True
    )
    # By default, get the most recent one:
    first_ckpt = recoverer.find_checkpoint()
    recoverer.save_and_keep_only(
        meta={"error": 5}, min_keys=["error"], keep_recent=True
    )
    second_ckpt = recoverer.find_checkpoint()
    assert first_ckpt.meta["unixtime"] < second_ckpt.meta["unixtime"]
    recoverer.save_and_keep_only(
        meta={"error": 6}, min_keys=["error"], keep_recent=True
    )
    third_ckpt = recoverer.find_checkpoint()
    remaining_ckpts = recoverer.list_checkpoints()
    assert first_ckpt not in remaining_ckpts
    assert second_ckpt in remaining_ckpts
    assert third_ckpt in remaining_ckpts

    # With equal importance criteria, the latest checkpoint should always be
    # returned
    fourth_ckpt = recoverer.save_checkpoint(meta={"error": 5})
    found_ckpt = recoverer.find_checkpoint(min_key="error")
    assert found_ckpt == fourth_ckpt
    fifth_ckpt = recoverer.save_checkpoint(meta={"error": 5})
    # Similarly for getting multiple checkpoints:
    found_ckpts = recoverer.find_checkpoints(
        min_key="error", max_num_checkpoints=2
    )
    assert found_ckpts == [fifth_ckpt, fourth_ckpt]


def test_torch_meta(tmpdir, device):
    from speechbrain.utils.checkpoints import Checkpointer
    import torch

    class Recoverable(torch.nn.Module):
        def __init__(self, param):
            super().__init__()
            self.param = torch.nn.Parameter(
                torch.tensor([param], device=device)
            )

        def forward(self, x):
            return x * self.param

    recoverable = Recoverable(1.0)
    recoverables = {"recoverable": recoverable}
    recoverer = Checkpointer(tmpdir, recoverables)
    saved = recoverer.save_checkpoint(
        meta={"loss": torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device)}
    )
    loaded = recoverer.recover_if_possible()
    assert saved.meta["loss"].allclose(loaded.meta["loss"])


def test_checkpoint_hook_register(tmpdir):
    from speechbrain.utils.checkpoints import register_checkpoint_hooks
    from speechbrain.utils.checkpoints import mark_as_saver
    from speechbrain.utils.checkpoints import mark_as_loader
    from speechbrain.utils.checkpoints import Checkpointer

    # First a proper interface:
    @register_checkpoint_hooks
    class CustomRecoverable:
        def __init__(self, param):
            self.param = int(param)

        @mark_as_saver
        def save(self, path):
            with open(path, "w") as fo:
                fo.write(str(self.param))

        @mark_as_loader
        def load(self, path, end_of_epoch, device):
            del end_of_epoch  # Unused
            with open(path) as fi:
                self.param = int(fi.read())

    recoverable = CustomRecoverable(1.0)
    checkpointer = Checkpointer(tmpdir, {"recoverable": recoverable})
    checkpointer.save_checkpoint()
    recoverable.param = 2.0
    checkpointer.recover_if_possible()
    assert recoverable.param == 1.0

    # Improper interfaces:
    with pytest.raises(TypeError):

        class BadRecoverable:
            def __init__(self, param):
                self.param = int(param)

            def save(self, path):
                with open(path, "w") as fo:
                    fo.write(str(self.param))

            @mark_as_loader
            def load(self, path, end_of_epoch):  # MISSING device
                del end_of_epoch  # Unused
                with open(path) as fi:
                    self.param = int(fi.read())

    with pytest.raises(TypeError):

        class BadRecoverable:  # noqa: F811
            def __init__(self, param):
                self.param = int(param)

            @mark_as_saver
            def save(self, path, extra_arg):  # Extra argument
                with open(path, "w") as fo:
                    fo.write(str(self.param))

            def load(self, path, end_of_epoch, device):
                del end_of_epoch  # Unused
                with open(path) as fi:
                    self.param = int(fi.read())


def test_torch_defaults(tmpdir, device):
    from speechbrain.utils.checkpoints import Checkpointer
    import torch

    module = torch.nn.Linear(10, 10).to(device)
    optimizer = torch.optim.Adam(module.parameters())
    lr_scheduler = torch.optim.lr_scheduler.CyclicLR(
        optimizer, 0.1, 1.0, cycle_momentum=False
    )
    # ReduceLROnPlateau is on an _LRScheduler for some reason, so have a separate test for it
    another_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
    checkpointer = Checkpointer(
        tmpdir,
        recoverables={
            "module": module,
            "optimizer": optimizer,
            "scheduler": lr_scheduler,
            "scheduler2": another_scheduler,
        },
    )
    ckpt = checkpointer.save_checkpoint()
    # test the module:
    inp = torch.randn((3, 10), device=device)
    prev_output = module(inp)

    # Re-initialize everything
    module = torch.nn.Linear(10, 10, device=device)
    optimizer = torch.optim.Adam(module.parameters())
    lr_scheduler = torch.optim.lr_scheduler.CyclicLR(
        optimizer, 0.1, 1.0, cycle_momentum=False
    )
    another_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
    checkpointer = Checkpointer(
        tmpdir,
        recoverables={
            "module": module,
            "optimizer": optimizer,
            "scheduler": lr_scheduler,
            "scheduler2": another_scheduler,
        },
    )
    checkpointer.load_checkpoint(ckpt)
    assert torch.allclose(module(inp), prev_output)

tests/unittests/test_checkpoints.py
TzuchengChang/NASS
def test_parse_arguments():
    from speechbrain.core import parse_arguments

    filename, run_opts, overrides = parse_arguments(
        ["params.yaml", "--device=cpu", "--seed=3", "--data_folder", "TIMIT"]
    )
    assert filename == "params.yaml"
    assert run_opts["device"] == "cpu"
    assert overrides == "seed: 3\ndata_folder: TIMIT"


def test_brain(device):
    import torch
    from speechbrain.core import Brain, Stage
    from torch.optim import SGD

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)

    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )

    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)

    start_output = brain.compute_forward(inputs, Stage.VALID)
    start_loss = brain.compute_objectives(start_output, targets, Stage.VALID)
    brain.fit(epoch_counter=range(10), train_set=train_set, valid_set=valid_set)
    end_output = brain.compute_forward(inputs, Stage.VALID)
    end_loss = brain.compute_objectives(end_output, targets, Stage.VALID)
    assert end_loss < start_loss

tests/unittests/test_core.py
TzuchengChang/NASS
import torch
import os


def test_read_audio_info(tmpdir, device):
    from speechbrain.dataio.dataio import read_audio_info, write_audio

    test_waveform = torch.rand(32000, device=device)

    for ext in ["wav", "ogg", "mp3"]:
        audio_path = os.path.join(tmpdir, f"test.{ext}")
        write_audio(audio_path, test_waveform.cpu(), 16000)
        info = read_audio_info(audio_path)

        # NOTE: This wide check is introduced for codecs that cannot handle
        # exact frame counts, such as mp3, which appears to operate over chunks
        # of 384 frames.
        # This remains an exact read of the metadata, however, which is what we
        # want to test and is the reason for `read_audio_info` to exist (see
        # comments there).
        assert (
            32000 <= info.num_frames <= 35000
        ), f"expected consistent len for codec {ext}"
        assert info.sample_rate == 16000


def test_read_audio(tmpdir, device):
    from speechbrain.dataio.dataio import read_audio, write_audio

    test_waveform = torch.rand(16000, device=device)
    wavfile = os.path.join(tmpdir, "wave.wav")
    write_audio(wavfile, test_waveform.cpu(), 16000)

    # dummy annotation
    for i in range(3):
        start = torch.randint(0, 8000, (1,), device=device).item()
        stop = start + torch.randint(500, 1000, (1,), device=device).item()

        loaded_range = read_audio(
            {"file": wavfile, "start": start, "stop": stop}
        ).to(device)
        assert loaded_range.allclose(test_waveform[start:stop], atol=1e-4)

        loaded_omit_start = read_audio({"file": wavfile, "stop": stop}).to(
            device
        )
        assert loaded_omit_start.allclose(test_waveform[:stop], atol=1e-4)

        loaded_omit_stop = read_audio({"file": wavfile, "start": start}).to(
            device
        )
        assert loaded_omit_stop.allclose(test_waveform[start:], atol=1e-4)

        loaded_simple = read_audio(wavfile).to(device)
        assert loaded_simple.allclose(test_waveform, atol=1e-4)

        # set to equal when switching to the sox_io backend
        # assert torch.all(torch.eq(loaded, test_waveform[0, start:stop]))


def test_read_audio_multichannel(tmpdir, device):
    from speechbrain.dataio.dataio import read_audio_multichannel, write_audio

    test_waveform = torch.rand(16000, 2, device=device)
    wavfile = os.path.join(tmpdir, "wave.wav")
    # sf.write(wavfile, test_waveform, 16000, subtype="float")
    write_audio(wavfile, test_waveform.cpu(), 16000)

    # dummy annotation we save and load one multichannel file
    for i in range(2):
        start = torch.randint(0, 8000, (1,), device=device).item()
        stop = start + torch.randint(500, 1000, (1,), device=device).item()

        wav_obj = {"wav": {"files": [wavfile], "start": start, "stop": stop}}

        loaded = read_audio_multichannel(wav_obj["wav"]).to(device)
        assert loaded.allclose(test_waveform[start:stop, :], atol=1e-4)
        # set to equal when switching to the sox_io backend
        # assert torch.all(torch.eq(loaded, test_waveform[:,start:stop]))

    # we test now multiple files loading
    test_waveform_2 = torch.rand(16000, 2, device=device)
    wavfile_2 = os.path.join(tmpdir, "wave_2.wav")
    write_audio(wavfile_2, test_waveform_2.cpu(), 16000)
    # sf.write(wavfile_2, test_waveform_2, 16000, subtype="float")

    for i in range(2):
        start = torch.randint(0, 8000, (1,), device=device).item()
        stop = start + torch.randint(500, 1000, (1,), device=device).item()

        wav_obj = {
            "wav": {"files": [wavfile, wavfile_2], "start": start, "stop": stop}
        }

        loaded = read_audio_multichannel(wav_obj["wav"]).to(device)
        test_waveform3 = torch.cat(
            (test_waveform[start:stop, :], test_waveform_2[start:stop, :]), 1
        )
        assert loaded.allclose(test_waveform3, atol=1e-4)
        # set to equal when switching to the sox_io backend
        # assert torch.all(
        #     torch.eq(
        #         loaded,
        #         torch.cat(
        #             (test_waveform[:,start:stop], test_waveform_2[:,start:stop]), 0
        #         ),
        #     )
        # )

tests/unittests/test_data_io.py
TzuchengChang/NASS
import torch
import pytest


def test_saveable_dataloader(tmpdir, device):
    from speechbrain.dataio.dataloader import SaveableDataLoader

    save_file = tmpdir + "/dataloader.ckpt"
    dataset = torch.randn(10, 1, device=device)
    dataloader = SaveableDataLoader(dataset, collate_fn=None)
    data_iterator = iter(dataloader)
    first_item = next(data_iterator)
    assert first_item == dataset[0]
    # Save here:
    dataloader._speechbrain_save(save_file)
    second_item = next(data_iterator)
    assert second_item == dataset[1]
    # Now make a new dataloader and recover:
    new_dataloader = SaveableDataLoader(dataset, collate_fn=None)
    new_dataloader._speechbrain_load(save_file, end_of_epoch=False, device=None)
    new_data_iterator = iter(new_dataloader)
    second_second_item = next(new_data_iterator)
    assert second_second_item == second_item


def test_saveable_dataloader_multiprocess(tmpdir):
    # Same test as above, but with multiprocess dataloading
    from speechbrain.dataio.dataloader import SaveableDataLoader

    save_file = tmpdir + "/dataloader.ckpt"
    dataset = torch.randn(10, 1)
    for num_parallel in [1, 2, 3, 4]:
        dataloader = SaveableDataLoader(
            dataset, num_workers=num_parallel, collate_fn=None
        )  # Note num_workers
        data_iterator = iter(dataloader)
        first_item = next(data_iterator)
        assert first_item == dataset[0]
        # Save here, note that this overwrites.
        dataloader._speechbrain_save(save_file)
        second_item = next(data_iterator)
        assert second_item == dataset[1]
        # Cleanup needed for MacOS (open file limit)
        del data_iterator
        del dataloader
        # Now make a new dataloader and recover:
        new_dataloader = SaveableDataLoader(
            dataset, num_workers=num_parallel, collate_fn=None
        )
        new_dataloader._speechbrain_load(
            save_file, end_of_epoch=False, device=None
        )
        new_data_iterator = iter(new_dataloader)
        second_second_item = next(new_data_iterator)
        assert second_second_item == second_item
        del new_data_iterator
        del new_dataloader


def test_looped_loader(tmpdir):
    # Tests that LoopedLoader will raise StopIteration appropriately
    # And that it can recover and keep the place.
    from speechbrain.dataio.dataloader import LoopedLoader

    save_file = tmpdir + "/loopedloader.ckpt"
    data = range(3)
    dataloader = LoopedLoader(data, epoch_length=2)
    data_iterator = iter(dataloader)
    assert next(data_iterator) == 0
    # Save here, 1 to go:
    dataloader.save(save_file)
    assert next(data_iterator) == 1
    with pytest.raises(StopIteration):
        next(data_iterator)
    # And it can be continued past the range:
    assert next(data_iterator) == 2
    assert next(data_iterator) == 0
    # And again it raises:
    with pytest.raises(StopIteration):
        next(data_iterator)
    # Now make a new dataloader and recover:
    new_dataloader = LoopedLoader(data, epoch_length=2)
    new_dataloader.load(save_file, end_of_epoch=False, device=None)
    new_data_iterator = iter(new_dataloader)
    next(new_data_iterator)
    with pytest.raises(StopIteration):
        next(new_data_iterator)

tests/unittests/test_dataloader.py
TzuchengChang/NASS
import torch
import torch.nn


def test_dropout(device):

    from speechbrain.nnet.dropout import Dropout2d

    inputs = torch.rand([4, 10, 32], device=device)
    drop = Dropout2d(drop_rate=0.0).to(device)
    outputs = drop(inputs)
    assert torch.all(torch.eq(inputs, outputs))

    drop = Dropout2d(drop_rate=1.0).to(device)
    outputs = drop(inputs)
    assert torch.all(
        torch.eq(torch.zeros(inputs.shape, device=device), outputs)
    )

    assert torch.jit.trace(drop, inputs)

tests/unittests/test_dropout.py
TzuchengChang/NASS
import torch


def test_embedding(device):

    from speechbrain.nnet.embedding import Embedding

    # create one hot vector and consider blank as zero vector
    embedding_dim = 39
    blank_id = 39
    size_dict = 40
    emb = Embedding(
        num_embeddings=size_dict, consider_as_one_hot=True, blank_id=blank_id
    ).to(device)
    inputs = torch.Tensor([10, 5, 2, 0, 39]).to(device).long()
    output = emb(inputs)
    assert output.shape == (5, 39)

    # use standard embedding layer
    embedding_dim = 128
    emb = Embedding(num_embeddings=size_dict, embedding_dim=embedding_dim).to(
        device
    )
    inputs = torch.randint(0, 40, (5, 10), device=device)
    output = emb(inputs)
    assert output.shape == (5, 10, 128)

    assert torch.jit.trace(emb, inputs)

tests/unittests/test_embedding.py
TzuchengChang/NASS
import torch


def test_deltas(device):

    from speechbrain.processing.features import Deltas

    size = torch.Size([10, 101, 20], device=device)
    inp = torch.ones(size, device=device)
    compute_deltas = Deltas(input_size=20).to(device)
    out = torch.zeros(size, device=device)
    assert torch.sum(compute_deltas(inp) == out) == out.numel()

    assert torch.jit.trace(compute_deltas, inp)


def test_context_window(device):

    from speechbrain.processing.features import ContextWindow

    inp = (
        torch.tensor([1, 2, 3], device=device)
        .unsqueeze(0)
        .unsqueeze(-1)
        .float()
    )
    compute_cw = ContextWindow(left_frames=1, right_frames=1).to(device)
    out = (
        torch.tensor([[0, 1, 2], [1, 2, 3], [2, 3, 0]], device=device)
        .unsqueeze(0)
        .float()
    )
    assert torch.sum(compute_cw(inp) == out) == 9

    inp = torch.rand([2, 10, 5], device=device)
    compute_cw = ContextWindow(left_frames=0, right_frames=0).to(device)
    assert torch.sum(compute_cw(inp) == inp) == inp.numel()

    assert torch.jit.trace(compute_cw, inp)


def test_istft(device):
    from speechbrain.processing.features import STFT
    from speechbrain.processing.features import ISTFT

    fs = 16000
    inp = torch.randn([10, 16000], device=device)
    inp = torch.stack(3 * [inp], -1)

    compute_stft = STFT(sample_rate=fs).to(device)
    compute_istft = ISTFT(sample_rate=fs).to(device)
    out = compute_istft(compute_stft(inp), sig_length=16000)

    assert torch.sum(torch.abs(inp - out) < 5e-5) >= inp.numel() - 5

    assert torch.jit.trace(compute_stft, inp)
    assert torch.jit.trace(compute_istft, compute_stft(inp))


def test_filterbank(device):

    from speechbrain.processing.features import Filterbank

    compute_fbanks = Filterbank().to(device)
    inputs = torch.ones([10, 101, 201], device=device)
    assert torch.jit.trace(compute_fbanks, inputs)

    # Check amin (-100 dB)
    inputs = torch.zeros([10, 101, 201], device=device)
    fbanks = compute_fbanks(inputs)
    assert torch.equal(fbanks, torch.ones_like(fbanks) * -100)

    # Check top_db
    fbanks = torch.zeros([1, 1, 1], device=device)
    expected = torch.Tensor([[[-100]]]).to(device)
    fbanks_db = compute_fbanks._amplitude_to_DB(fbanks)
    assert torch.equal(fbanks_db, expected)

    # Making sure independent computation gives same results
    # as the batch computation
    input1 = torch.rand([1, 101, 201], device=device) * 10
    input2 = torch.rand([1, 101, 201], device=device)
    input3 = torch.cat([input1, input2], dim=0)
    fbank1 = compute_fbanks(input1)
    fbank2 = compute_fbanks(input2)
    fbank3 = compute_fbanks(input3)
    assert torch.sum(torch.abs(fbank1[0] - fbank3[0])) < 8e-05
    assert torch.sum(torch.abs(fbank2[0] - fbank3[1])) < 8e-05


def test_dtc(device):

    from speechbrain.processing.features import DCT

    compute_dct = DCT(input_size=40)
    inputs = torch.randn([10, 101, 40], device=device)
    assert torch.jit.trace(compute_dct, inputs)


def test_input_normalization(device):

    from speechbrain.processing.features import InputNormalization

    norm = InputNormalization().to(device)
    inputs = torch.randn([10, 101, 20], device=device)
    inp_len = torch.ones([10], device=device)
    assert torch.jit.trace(norm, (inputs, inp_len))

    norm = InputNormalization().to(device)
    inputs = (
        torch.FloatTensor([1, 2, 3, 0, 0, 0])
        .to(device)
        .unsqueeze(0)
        .unsqueeze(2)
    )
    inp_len = torch.FloatTensor([0.5]).to(device)
    out_norm = norm(inputs, inp_len).squeeze()
    target = torch.FloatTensor([-1, 0, 1, -2, -2, -2]).to(device)
    assert torch.equal(out_norm, target)


def test_features_multimic(device):

    from speechbrain.processing.features import Filterbank

    compute_fbanks = Filterbank().to(device)
    inputs = torch.rand([10, 101, 201], device=device)
    output = compute_fbanks(inputs)
    inputs_ch2 = torch.stack((inputs, inputs), -1)
    output_ch2 = compute_fbanks(inputs_ch2)
    output_ch2 = output_ch2[..., 0]
    assert torch.sum(output - output_ch2) < 5e-05

tests/unittests/test_features.py
TzuchengChang/NASS
import torch
from torch.nn import functional as F


def _fake_probs(idx, count):
    result = torch.zeros(count)
    result[idx] = 2.0
    return F.softmax(result, dim=-1)


def _batch_fake_probs(indexes, count):
    p_seq = torch.zeros(indexes.shape + (count,))

    for batch_idx in range(len(indexes)):
        for item_idx in range(indexes.size(1)):
            p_seq[batch_idx, item_idx, :] = _fake_probs(
                indexes[batch_idx, item_idx], count
            )
    return p_seq


def test_subsequence_loss():
    from speechbrain.nnet.losses import nll_loss
    from speechbrain.lobes.models.g2p.homograph import SubsequenceLoss

    phn_dim = 4
    phns = torch.tensor(
        [
            [1, 2, 3, 0, 3, 1, 2, 1, 0, 3, 1, 2, 0, 0],
            [1, 2, 3, 1, 0, 3, 2, 1, 0, 1, 3, 2, 0, 0],
            [1, 2, 3, 1, 2, 3, 0, 1, 3, 1, 3, 2, 0, 1],
        ]
    )
    phn_lens = torch.tensor([12, 12, 14])

    preds = torch.tensor(
        [
            [1, 3, 3, 0, 3, 1, 2, 1, 0, 3, 1, 2],
            [1, 1, 2, 1, 0, 3, 2, 1, 0, 1, 3, 2],
            [3, 2, 1, 1, 2, 3, 0, 1, 3, 2, 3, 3],
        ]
    )

    p_seq = _batch_fake_probs(preds, phn_dim)

    start = torch.tensor([0, 5, 7])
    end = torch.tensor([3, 8, 12])

    word_phns_pred = torch.tensor(
        [[1, 3, 3, 0, 0], [3, 2, 1, 0, 0], [1, 3, 2, 3, 3]]
    )
    word_phns_ref = torch.tensor(
        [[1, 2, 3, 0, 0], [3, 2, 1, 0, 0], [1, 3, 1, 3, 2]]
    )
    word_p_seq = _batch_fake_probs(word_phns_pred, phn_dim)
    word_lengths = torch.tensor([3, 3, 5]) / 5

    loss = SubsequenceLoss(seq_cost=nll_loss, word_separator=0)
    loss_value = loss.forward(phns, phn_lens, p_seq.log(), start, end)
    loss_value_ref = nll_loss(word_p_seq.log(), word_phns_ref, word_lengths)
    assert loss_value == loss_value_ref


def test_extract_hyps():
    from speechbrain.lobes.models.g2p.homograph import SubsequenceExtractor

    phns = torch.tensor(
        [
            [1, 2, 3, 0, 3, 1, 2, 1, 0, 3, 1, 2, 0, 0],
            [1, 2, 3, 1, 0, 3, 2, 1, 0, 1, 3, 2, 0, 0],
            [1, 2, 3, 1, 2, 3, 0, 1, 3, 1, 3, 2, 0, 1],
        ]
    )
    hyps = [
        [1, 2, 3, 2, 0, 3, 1, 2, 1, 0, 3, 1, 2],
        [1, 2, 3, 0, 3, 2, 1, 0, 1, 3, 2],
        [1, 2, 3, 1, 2, 3, 0, 1, 3, 1, 3, 2, 0, 1],
    ]
    subsequence_phn_start = torch.tensor([4, 0, 7])
    ref_hyps = [[3, 1, 2, 1], [1, 2, 3], [1, 3, 1, 3, 2]]

    extractor = SubsequenceExtractor(word_separator=0)
    subsequence_hyps = extractor.extract_hyps(
        ref_seq=phns, hyps=hyps, subsequence_phn_start=subsequence_phn_start
    )
    assert subsequence_hyps == ref_hyps

tests/unittests/test_g2p.py
TzuchengChang/NASS
import pytest


def test_hpopt_generic():
    from io import StringIO
    from speechbrain.utils import hpopt as hp
    import json

    output = StringIO()

    reporter = hp.GenericHyperparameterOptimizationReporter(
        objective_key="per", output=output
    )
    result = {"train_loss": 0.9, "valid_loss": 1.2, "per": 0.10}
    reporter.report_objective(result)
    output.seek(0)
    output_result = json.load(output)
    assert output_result["train_loss"] == pytest.approx(0.9)
    assert output_result["valid_loss"] == pytest.approx(1.2)
    assert output_result["per"] == pytest.approx(0.10)
    assert output_result["objective"] == pytest.approx(0.10)


def test_hpopt_orion():
    from speechbrain.utils import hpopt as hp

    results = {}

    class MockOrion:
        def report_objective(self, value):
            results["value"] = value

    mock_orion = MockOrion()

    reporter = hp.OrionHyperparameterOptimizationReporter(
        objective_key="valid_loss"
    )
    reporter.orion_client = mock_orion

    result = {"train_loss": 0.9, "valid_loss": 1.2, "per": 0.10}
    reporter.report_objective(result)
    assert results["value"] == pytest.approx(1.2)


def test_hpopt_context():
    import json
    from speechbrain.utils import hpopt as hp
    from io import StringIO

    output = StringIO()
    reporter = hp.GenericHyperparameterOptimizationReporter(
        objective_key="per", output=output
    )

    with hp.hyperparameter_optimization() as hp_ctx:
        hp_ctx.reporter = reporter
        result = {"per": 10, "loss": 1.2}
        hp.report_result(result)

        result = {"per": 3, "loss": 1.3}
        hp.report_result(result)

    output.seek(0)
    output_result = json.load(output)
    assert output_result["per"] == 3

tests/unittests/test_hpopt.py
TzuchengChang/NASS
import torch
import torch.nn


def test_linear(device):

    from speechbrain.nnet.linear import Linear

    inputs = torch.rand(1, 2, 4, device=device)
    lin_t = Linear(n_neurons=4, input_size=inputs.shape[-1], bias=False)
    lin_t.w.weight = torch.nn.Parameter(
        torch.eye(inputs.shape[-1], device=device)
    )
    outputs = lin_t(inputs)
    assert torch.all(torch.eq(inputs, outputs))

    assert torch.jit.trace(lin_t, inputs)

tests/unittests/test_linear.py
TzuchengChang/NASS
import torch
import pytest


def test_nll(device):
    from speechbrain.nnet.losses import nll_loss

    predictions = torch.zeros(4, 10, 8, device=device)
    targets = torch.zeros(4, 10, device=device)
    lengths = torch.ones(4, device=device)
    out_cost = nll_loss(predictions, targets, lengths)
    assert torch.all(torch.eq(out_cost, 0))


def test_mse(device):
    from speechbrain.nnet.losses import mse_loss

    predictions = torch.ones(4, 10, 8, device=device)
    targets = torch.ones(4, 10, 8, device=device)
    lengths = torch.ones(4, device=device)
    out_cost = mse_loss(predictions, targets, lengths)
    assert torch.all(torch.eq(out_cost, 0))

    predictions = torch.zeros(4, 10, 8, device=device)
    out_cost = mse_loss(predictions, targets, lengths)
    assert torch.all(torch.eq(out_cost, 1))


def test_l1(device):
    from speechbrain.nnet.losses import l1_loss

    predictions = torch.ones(4, 10, 8, device=device)
    targets = torch.ones(4, 10, 8, device=device)
    lengths = torch.ones(4, device=device)
    out_cost = l1_loss(predictions, targets, lengths)
    assert torch.all(torch.eq(out_cost, 0))


def test_bce_loss(device):
    from speechbrain.nnet.losses import bce_loss

    # Ensure this works both with and without singleton dimension
    predictions_singleton = torch.zeros(4, 10, 1, device=device)
    predictions_match = torch.zeros(4, 10, device=device)
    targets = torch.ones(4, 10, device=device)
    lengths = torch.ones(4, device=device)
    out_cost_singleton = bce_loss(predictions_singleton, targets, lengths)
    out_cost_match = bce_loss(predictions_match, targets, lengths)
    assert torch.allclose(
        torch.exp(out_cost_singleton), torch.tensor(2.0, device=device)
    )
    assert torch.allclose(
        torch.exp(out_cost_match), torch.tensor(2.0, device=device)
    )

    # How about one dimensional inputs
    predictions = torch.zeros(5, 1, device=device)
    targets = torch.ones(5, device=device)
    out_cost = bce_loss(predictions, targets)
    assert torch.allclose(torch.exp(out_cost), torch.tensor(2.0, device=device))

    # Can't pass lengths in 1D case
    with pytest.raises(ValueError):
        bce_loss(predictions, targets, length=torch.ones(5, device=device))


def test_classification_error(device):
    from speechbrain.nnet.losses import classification_error

    predictions = torch.zeros(4, 10, 8, device=device)
    predictions[:, :, 0] += 1.0
    targets = torch.zeros(4, 10, device=device)
    lengths = torch.ones(4, device=device)
    out_cost = classification_error(predictions, targets, lengths)
    assert torch.all(torch.eq(out_cost, 0))


def test_pitwrapper(device):
    from speechbrain.nnet.losses import PitWrapper
    import torch
    from torch import nn

    base_loss = nn.MSELoss(reduction="none")
    pit = PitWrapper(base_loss)
    predictions = torch.rand(
        (2, 32, 4), device=device
    )  # batch, frames, sources
    p = (3, 0, 2, 1)
    # same but we invert the ordering to check if permutation invariant
    targets = predictions[..., p]
    loss, opt_p = pit(predictions, targets)
    assert [x == p for x in opt_p] == [True for i in range(len(opt_p))]
    predictions = pit.reorder_tensor(predictions, opt_p)
    assert torch.all(torch.eq(base_loss(predictions, targets), 0))

    predictions = torch.rand(
        (3, 32, 32, 32, 5), device=device
    )  # batch, ..., sources
    p = (3, 0, 2, 1, 4)
    targets = predictions[
        ..., p
    ]  # same but we invert the ordering to check if permutation invariant
    loss, opt_p = pit(predictions, targets)
    assert [x == p for x in opt_p] == [True for i in range(len(opt_p))]
    predictions = pit.reorder_tensor(predictions, opt_p)
    assert torch.all(torch.eq(base_loss(predictions, targets), 0))


def test_transducer_loss(device):
    # Make this its own test since it can only be run
    # if numba is installed and a GPU is available
    pytest.importorskip("numba")
    if torch.cuda.device_count() == 0:
        pytest.skip("This test can only be run if a GPU is available")

    from speechbrain.nnet.losses import transducer_loss

    device = torch.device("cuda")
    log_probs = (
        torch.Tensor(
            [
                [
                    [
                        [0.1, 0.6, 0.1, 0.1, 0.1],
                        [0.1, 0.1, 0.6, 0.1, 0.1],
                        [0.1, 0.1, 0.2, 0.8, 0.1],
                    ],
                    [
                        [0.1, 0.6, 0.1, 0.1, 0.1],
                        [0.1, 0.1, 0.2, 0.1, 0.1],
                        [0.7, 0.1, 0.2, 0.1, 0.1],
                    ],
                ]
            ]
        )
        .to(device)
        .requires_grad_()
        .log_softmax(dim=-1)
    )
    targets = torch.Tensor([[1, 2]]).to(device).int()
    probs_length = torch.Tensor([1.0]).to(device)
    target_length = torch.Tensor([1.0]).to(device)
    out_cost = transducer_loss(
        log_probs,
        targets,
        probs_length,
        target_length,
        blank_index=0,
        use_torchaudio=False,
    )
    out_cost.backward()
    assert out_cost.item() == pytest.approx(2.2478, 0.0001)


def test_guided_attention_loss_mask(device):
    from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss

    loss = GuidedAttentionLoss().to(device)
    input_lengths = torch.tensor([3, 2, 6], device=device)
    output_lengths = torch.tensor([4, 3, 5], device=device)
    soft_mask = loss.guided_attentions(input_lengths, output_lengths)
    ref_soft_mask = torch.tensor(
        [
            [
                [0.0, 0.54216665, 0.9560631, 0.9991162, 0.0],
                [0.7506478, 0.08314464, 0.2933517, 0.8858382, 0.0],
                [0.9961341, 0.8858382, 0.2933517, 0.08314464, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ],
            [
                [0.0, 0.7506478, 0.9961341, 0.0, 0.0],
                [0.9560631, 0.2933517, 0.2933517, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ],
            [
                [0.0, 0.39346933, 0.86466473, 0.988891, 0.99966455],
                [0.2933517, 0.01379288, 0.49366438, 0.90436554, 0.993355],
                [0.7506478, 0.1992626, 0.05404053, 0.5888877, 0.93427145],
                [0.9560631, 0.6753475, 0.1175031, 0.1175031, 0.6753475],
                [0.9961341, 0.93427145, 0.5888877, 0.05404053, 0.1992626],
                [0.9998301, 0.993355, 0.90436554, 0.49366438, 0.01379288],
            ],
        ],
        device=device,
    )
    assert torch.allclose(soft_mask, ref_soft_mask)


def test_guided_attention_loss_value(device):
    from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss

    loss = GuidedAttentionLoss().to(device)
    input_lengths = torch.tensor([2, 3], device=device)
    target_lengths = torch.tensor([3, 4], device=device)
    alignments = torch.tensor(
        [
            [
                [0.8, 0.2, 0.0],
                [0.4, 0.6, 0.0],
                [0.2, 0.8, 0.0],
                [0.0, 0.0, 0.0],
            ],
            [
                [0.6, 0.2, 0.2],
                [0.1, 0.7, 0.2],
                [0.3, 0.4, 0.3],
                [0.2, 0.3, 0.5],
            ],
        ],
        device=device,
    )
    loss_value = loss(alignments, input_lengths, target_lengths)
    ref_loss_value = torch.tensor(0.1142)
    assert torch.isclose(loss_value, ref_loss_value, 0.0001, 0.0001).item()


def test_guided_attention_loss_shapes(device):
    from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss

    loss = GuidedAttentionLoss().to(device)
    input_lengths = torch.tensor([3, 2, 6], device=device)
    output_lengths = torch.tensor([4, 3, 5], device=device)
    soft_mask = loss.guided_attentions(input_lengths, output_lengths)
    assert soft_mask.shape == (3, 6, 5)
    soft_mask = loss.guided_attentions(
        input_lengths, output_lengths, max_input_len=10
    )
    assert soft_mask.shape == (3, 10, 5)
    soft_mask = loss.guided_attentions(
        input_lengths, output_lengths, max_target_len=12
    )
    assert soft_mask.shape == (3, 6, 12)
    soft_mask = loss.guided_attentions(
        input_lengths, output_lengths, max_input_len=10, max_target_len=12
    )
    assert soft_mask.shape == (3, 10, 12)

tests/unittests/test_losses.py
TzuchengChang/NASS
import torch
import torch.nn
import math


def test_metric_stats(device):
    from speechbrain.utils.metric_stats import MetricStats
    from speechbrain.nnet.losses import l1_loss

    l1_stats = MetricStats(metric=l1_loss)
    l1_stats.append(
        ids=["utterance1", "utterance2"],
        predictions=torch.tensor([[0.1, 0.2], [0.1, 0.2]], device=device),
        targets=torch.tensor([[0.1, 0.3], [0.2, 0.3]], device=device),
        length=torch.ones(2, device=device),
        reduction="batch",
    )
    summary = l1_stats.summarize()
    assert math.isclose(summary["average"], 0.075, rel_tol=1e-5)
    assert math.isclose(summary["min_score"], 0.05, rel_tol=1e-5)
    assert summary["min_id"] == "utterance1"
    assert math.isclose(summary["max_score"], 0.1, rel_tol=1e-5)
    assert summary["max_id"] == "utterance2"


def test_error_rate_stats(device):
    from speechbrain.utils.metric_stats import ErrorRateStats

    wer_stats = ErrorRateStats()
    i2l = {1: "hello", 2: "world", 3: "the"}

    def mapper(batch):
        return [[i2l[int(x)] for x in seq] for seq in batch]

    wer_stats.append(
        ids=["utterance1", "utterance2"],
        predict=[[3, 2, 1], [2, 3]],
        target=torch.tensor([[3, 2, 0], [2, 1, 0]], device=device),
        target_len=torch.tensor([0.67, 0.67], device=device),
        ind2lab=mapper,
    )
    summary = wer_stats.summarize()
    assert summary["WER"] == 50.0
    assert summary["insertions"] == 1
    assert summary["substitutions"] == 1
    assert summary["deletions"] == 0
    assert wer_stats.scores[0]["ref_tokens"] == ["the", "world"]
    assert wer_stats.scores[0]["hyp_tokens"] == ["the", "world", "hello"]


def test_binary_metrics(device):
    from speechbrain.utils.metric_stats import BinaryMetricStats

    binary_stats = BinaryMetricStats()
    binary_stats.append(
        ids=["utt1", "utt2", "utt3", "utt4", "utt5", "utt6"],
        scores=torch.tensor([0.1, 0.4, 0.8, 0.2, 0.3, 0.6], device=device),
        labels=torch.tensor([1, 0, 1, 0, 1, 0], device=device),
    )
    summary = binary_stats.summarize(threshold=0.5)
    assert summary["TP"] == 1
    assert summary["TN"] == 2
    assert summary["FP"] == 1
    assert summary["FN"] == 2

    summary = binary_stats.summarize(threshold=None)
    assert summary["threshold"] >= 0.3 and summary["threshold"] < 0.4

    summary = binary_stats.summarize(threshold=None, max_samples=1)
    assert summary["threshold"] >= 0.1 and summary["threshold"] < 0.2


def test_EER(device):
    from speechbrain.utils.metric_stats import EER

    positive_scores = torch.tensor([0.1, 0.2, 0.3], device=device)
    negative_scores = torch.tensor([0.4, 0.5, 0.6], device=device)
    eer, threshold = EER(positive_scores, negative_scores)
    assert eer == 1.0
    assert threshold > 0.3 and threshold < 0.4

    positive_scores = torch.tensor([0.4, 0.5, 0.6], device=device)
    negative_scores = torch.tensor([0.3, 0.2, 0.1], device=device)
    eer, threshold = EER(positive_scores, negative_scores)
    assert eer == 0
    assert threshold > 0.3 and threshold < 0.4

    cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
    input1 = torch.randn(1000, 64, device=device)
    input2 = torch.randn(1000, 64, device=device)
    positive_scores = cos(input1, input2)

    input1 = torch.randn(1000, 64, device=device)
    input2 = torch.randn(1000, 64, device=device)
    negative_scores = cos(input1, input2)

    eer, threshold = EER(positive_scores, negative_scores)

    correct = (positive_scores > threshold).nonzero(as_tuple=False).size(0) + (
        negative_scores < threshold
    ).nonzero(as_tuple=False).size(0)

    assert correct > 900 and correct < 1100


def test_minDCF(device):
    from speechbrain.utils.metric_stats import minDCF

    positive_scores = torch.tensor([0.1, 0.2, 0.3], device=device)
    negative_scores = torch.tensor([0.4, 0.5, 0.6], device=device)
    min_dcf, threshold = minDCF(positive_scores, negative_scores)
    assert (0.01 - min_dcf) < 1e-4
    assert threshold >= 0.6

    positive_scores = torch.tensor([0.4, 0.5, 0.6], device=device)
    negative_scores = torch.tensor([0.1, 0.2, 0.3], device=device)
    min_dcf, threshold = minDCF(positive_scores, negative_scores)
    assert min_dcf == 0
    assert threshold > 0.3 and threshold < 0.4


def test_classification_stats():
    import pytest
    from speechbrain.utils.metric_stats import ClassificationStats

    stats = ClassificationStats()
    stats.append(ids=["1", "2"], predictions=["B", "A"], targets=["B", "A"])
    stats.append(ids=["3", "4"], predictions=["A", "B"], targets=["B", "C"])

    summary = stats.summarize()
    assert pytest.approx(summary["accuracy"], 0.01) == 0.5
    classwise_accuracy = summary["classwise_accuracy"]
    assert pytest.approx(classwise_accuracy["A"]) == 1.0
    assert pytest.approx(classwise_accuracy["B"]) == 0.5
    assert pytest.approx(classwise_accuracy["C"]) == 0.0


def test_categorized_classification_stats():
    import pytest
    from speechbrain.utils.metric_stats import ClassificationStats

    stats = ClassificationStats()
    stats.append(
        ids=["1", "2"],
        predictions=["B", "A"],
        targets=["B", "A"],
        categories=["C1", "C2"],
    )
    stats.append(
        ids=["3", "4"],
        predictions=["A", "B"],
        targets=["B", "C"],
        categories=["C2", "C1"],
    )
    stats.append(
        ids=["5", "6"],
        predictions=["A", "C"],
        targets=["B", "C"],
        categories=["C2", "C1"],
    )

    summary = stats.summarize()
    assert pytest.approx(summary["accuracy"], 0.01) == 0.5
    classwise_accuracy = summary["classwise_accuracy"]
    assert pytest.approx(classwise_accuracy["C1", "B"]) == 1.0
    assert pytest.approx(classwise_accuracy["C1", "C"]) == 0.5
    assert pytest.approx(classwise_accuracy["C2", "A"]) == 1.0
    assert pytest.approx(classwise_accuracy["C2", "B"]) == 0.0


def test_classification_stats_report():
    from io import StringIO
    from speechbrain.utils.metric_stats import ClassificationStats

    stats = ClassificationStats()
    stats.append(ids=["1", "2"], predictions=["B", "A"], targets=["B", "A"])
    stats.append(ids=["3", "4"], predictions=["A", "B"], targets=["B", "C"])
    report_file = StringIO()
    stats.write_stats(report_file)
    report_file.seek(0)
    report = report_file.read()
    ref_report = """Overall Accuracy: 50%

Class-Wise Accuracy
-------------------
A: 1 / 1 (100.00%)
B: 1 / 2 (50.00%)
C: 0 / 1 (0.00%)

Confusion
---------
Target: A
  -> A: 1 / 1 (100.00%)
Target: B
  -> A: 1 / 2 (50.00%)
  -> B: 1 / 2 (50.00%)
Target: C
  -> B: 1 / 1 (100.00%)
"""
    assert report == ref_report

tests/unittests/test_metrics.py
TzuchengChang/NASS
import torch


def test_gccphat(device):

    from speechbrain.processing.features import STFT
    from speechbrain.processing.multi_mic import Covariance, GccPhat

    # Creating the test signal
    fs = 16000

    delay = 60

    sig = torch.randn([10, fs], device=device)
    sig_delayed = torch.cat(
        (torch.zeros([10, delay], device=device), sig[:, 0:-delay]), 1
    )

    xs = torch.stack((sig_delayed, sig), -1)

    stft = STFT(sample_rate=fs).to(device)
    Xs = stft(xs)

    # Computing the covariance matrix for GCC-PHAT
    cov = Covariance().to(device)
    gccphat = GccPhat().to(device)

    XXs = cov(Xs).to(device)
    tdoas = torch.abs(gccphat(XXs))

    n_valid_tdoas = torch.sum(torch.abs(tdoas[..., 1] - delay) < 1e-3)
    assert n_valid_tdoas == Xs.shape[0] * Xs.shape[1]
    assert torch.jit.trace(stft, xs)
    assert torch.jit.trace(cov, Xs)
    assert torch.jit.trace(gccphat, XXs)

tests/unittests/test_multi_mic.py
TzuchengChang/NASS
import torch
import torch.nn


def test_BatchNorm1d(device):

    from speechbrain.nnet.normalization import BatchNorm1d

    input = torch.randn(100, 10, device=device) + 2.0
    norm = BatchNorm1d(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=0).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=0).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(100, 20, 10, device=device) + 2.0
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=0).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=0).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    # Test with combined dimensions
    input = torch.randn(100, 10, 20, device=device) + 2.0
    norm = BatchNorm1d(input_shape=input.shape, combine_batch_time=True).to(
        device
    )
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=0).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=0).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(100, 40, 20, 30, device=device) + 2.0
    norm = BatchNorm1d(input_shape=input.shape, combine_batch_time=True).to(
        device
    )
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=0).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=0).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)


def test_BatchNorm2d(device):

    from speechbrain.nnet.normalization import BatchNorm2d

    input = torch.randn(100, 10, 4, 20, device=device) + 2.0
    norm = BatchNorm2d(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=0).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=0).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)


def test_LayerNorm(device):

    from speechbrain.nnet.normalization import LayerNorm

    input = torch.randn(4, 101, 256, device=device) + 2.0
    norm = LayerNorm(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=2).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=2).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(100, 101, 16, 32, device=device) + 2.0
    norm = LayerNorm(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=[2, 3]).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=[2, 3]).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)


def test_InstanceNorm1d(device):

    from speechbrain.nnet.normalization import InstanceNorm1d

    input = torch.randn(100, 10, 128, device=device) + 2.0
    norm = InstanceNorm1d(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=2).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=2).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)


def test_InstanceNorm2d(device):

    from speechbrain.nnet.normalization import InstanceNorm2d

    input = torch.randn(100, 10, 20, 2, device=device) + 2.0
    norm = InstanceNorm2d(input_shape=input.shape).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=[2, 3]).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=[2, 3]).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)


def test_GroupNorm(device):

    from speechbrain.nnet.normalization import GroupNorm

    input = torch.randn(4, 101, 256, device=device) + 2.0
    norm = GroupNorm(input_shape=input.shape, num_groups=256).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=2).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=2).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(4, 101, 256, device=device) + 2.0
    norm = GroupNorm(input_shape=input.shape, num_groups=128).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=2).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=2).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(100, 101, 16, 32, device=device) + 2.0
    norm = GroupNorm(input_shape=input.shape, num_groups=32).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=3).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=3).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    input = torch.randn(100, 101, 16, 32, device=device) + 2.0
    norm = GroupNorm(input_shape=input.shape, num_groups=8).to(device)
    output = norm(input)
    assert input.shape == output.shape

    current_mean = output.mean(dim=3).mean()
    assert torch.abs(current_mean) < 1e-06

    current_std = output.std(dim=3).mean()
    assert torch.abs(1.0 - current_std) < 0.01

    assert torch.jit.trace(norm, input)

tests/unittests/test_normalization.py
TzuchengChang/NASS
import pytest

from speechbrain.utils.parallel import parallel_map

small_test_input = [1, 2, 3, 4, 5, 6] * 5
small_test_expected = [2, 4, 6, 8, 10, 12] * 5


def small_test_func(x):
    return x * 2


def raise_error(x):
    raise ValueError("dummy error")


def test_parallel_map():
    # test different chunk sizes
    for test_chunk_size in [1, 2, 4, 8, 16, 32, 64, 128]:
        small_test_output = list(
            parallel_map(
                small_test_func,
                small_test_input,
                process_count=2,
                chunk_size=test_chunk_size,
            )
        )

        assert (
            small_test_output == small_test_expected
        ), f"chunk_size={test_chunk_size}"

    # test whether pbar off works properly
    small_test_output = list(
        parallel_map(small_test_func, small_test_input, progress_bar=False)
    )

    # test whether exceptions are forwarded properly
    with pytest.raises(ValueError):
        small_test_output = list(
            parallel_map(raise_error, small_test_input, progress_bar=False)
        )

    # test edge case: empty input
    assert list(parallel_map(small_test_func, [])) == []

    # trivial test for tqdm kwargs
    parallel_map(small_test_func, small_test_input, progress_bar_kwargs={})

tests/unittests/test_parallel.py
TzuchengChang/NASS
import torch
import torch.nn


def test_pooling1d(device):

    from speechbrain.nnet.pooling import Pooling1d

    input = (
        torch.tensor([1, 3, 2], device=device)
        .unsqueeze(0)
        .unsqueeze(-1)
        .float()
    )
    pool = Pooling1d("max", 3).to(device)
    output = pool(input)
    assert output == 3

    pool = Pooling1d("avg", 3).to(device)
    output = pool(input)
    assert output == 2

    assert torch.jit.trace(pool, input)


def test_pooling2d(device):

    from speechbrain.nnet.pooling import Pooling2d

    input = (
        torch.tensor([[1, 3, 2], [4, 6, 5]], device=device).float().unsqueeze(0)
    )
    pool = Pooling2d("max", (2, 3)).to(device)
    output = pool(input)
    assert output == 6

    input = (
        torch.tensor([[1, 3, 2], [4, 6, 5]], device=device).float().unsqueeze(0)
    )
    pool = Pooling2d("max", (1, 3)).to(device)
    output = pool(input)
    assert output[0][0] == 3
    assert output[0][1] == 6

    input = (
        torch.tensor([[1, 3, 2], [4, 6, 5]], device=device).float().unsqueeze(0)
    )
    pool = Pooling2d("avg", (2, 3)).to(device)
    output = pool(input)
    assert output == 3.5

    input = (
        torch.tensor([[1, 3, 2], [4, 6, 5]], device=device).float().unsqueeze(0)
    )
    pool = Pooling2d("avg", (1, 3)).to(device)
    output = pool(input)
    assert output[0][0] == 2
    assert output[0][1] == 5

    assert torch.jit.trace(pool, input)

tests/unittests/test_pooling.py
TzuchengChang/NASS
def test_pretrainer(tmpdir, device):
    import torch
    from torch.nn import Linear

    # save a model in tmpdir/original/model.ckpt
    first_model = Linear(32, 32).to(device)
    pretrained_dir = tmpdir / "original"
    pretrained_dir.mkdir()
    with open(pretrained_dir / "model.ckpt", "wb") as fo:
        torch.save(first_model.state_dict(), fo)

    # Make a new model and Pretrainer
    pretrained_model = Linear(32, 32).to(device)
    assert not torch.all(torch.eq(pretrained_model.weight, first_model.weight))
    from speechbrain.utils.parameter_transfer import Pretrainer

    pt = Pretrainer(
        collect_in=tmpdir / "reused", loadables={"model": pretrained_model}
    )
    pt.collect_files(default_source=pretrained_dir)
    pt.load_collected()
    assert torch.all(torch.eq(pretrained_model.weight, first_model.weight))

tests/unittests/test_pretrainer.py
TzuchengChang/NASS
def test_profile_class(device):
    import torch
    from torch.optim import SGD
    from speechbrain.core import Brain
    from speechbrain.utils.profiling import profile

    @profile
    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)

    # Profiling: __init__ constructor.
    brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    assert brain.profiler is not None
    assert brain.profiler.profiler is not None
    # assert len(brain.profiler.key_averages()) == 2
    # assert (brain.profiler.events().total_average().count >= 4)  # == 6  # before; config dependent: 7
    assert (
        len(brain.profiler.speechbrain_event_traces) == 1
    )  # set & filled by the @profile decorator
    """print(brain.profiler.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                           Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
        aten::_has_compatible_shallow_copy_type        73.33%      11.000us        73.33%      11.000us       2.750us             4
                                       aten::to        26.67%       4.000us        26.67%       4.000us       2.000us             2
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 15.000us
    """

    # Profiling: fit() for train operations.
    # By default, @profile should also annotate fit & evaluate functions; here the fit function is tested only.
    brain.fit(epoch_counter=range(10), train_set=train_set, valid_set=valid_set)
    assert brain.profiler is not None
    # assert len(brain.profiler.key_averages()) >= 60  # 72 with torch==1.10.1
    # assert brain.profiler.events().total_average().count >= 2000  # 2832 with torch==1.10.1
    assert len(brain.profiler.speechbrain_event_traces) == 2
    # assert len(brain.profiler.speechbrain_event_traces[0]) >= 4  # == 6  # before; config dependent: 7
    # assert len(brain.profiler.speechbrain_event_traces[1]) >= 2000  # 2862 with torch==1.10.1
    """print(brain.profiler.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              aten::l1_loss         2.60%     443.000us        26.15%       4.460ms     111.500us            40
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        12.97%       2.212ms        20.28%       3.459ms      86.475us            40
                                               aten::linear         1.28%     219.000us        12.97%       2.212ms     110.600us            20
                                                   aten::to         1.88%     320.000us        10.68%       1.822ms      11.387us           160
                                             aten::isfinite         1.98%     338.000us         9.82%       1.674ms      55.800us            30
                                                 aten::mean         1.10%     188.000us         9.31%       1.587ms      79.350us            20
                                             aten::_to_copy         5.65%     964.000us         8.99%       1.533ms      15.330us           100
                                                aten::stack         2.23%     380.000us         8.67%       1.479ms      29.580us            50
                                               aten::matmul         1.62%     277.000us         7.63%       1.301ms      65.050us            20
                                     aten::l1_loss_backward         1.24%     212.000us         6.78%       1.157ms      57.850us            20
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 17.054ms
    """


def test_profile_func(device):
    # import torch
    # from pytest import raises
    # from torch.optim import SGD
    # from speechbrain.core import Brain
    # from torch.autograd.profiler import record_function
    from speechbrain.utils.profiling import profile

    # from speechbrain.utils.profiling import events_diff

    """
    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    class SimpleBrainNittyGritty(Brain):
        def compute_forward(self, batch, stage):
            # example: one way of using torch.autograd.profiler.record_function
            with record_function("is this faster (?)"):
                this = self.modules.model(batch[0])
            return this

        def compute_objectives(self, predictions, batch, stage):
            # example: one could also think of running comparative testing using record_function
            with record_function("or that (?)"):
                that = torch.nn.functional.l1_loss(predictions, batch[1])
            return that
    """

    @profile
    def train(brain, train_set, valid_set):
        brain.fit(
            epoch_counter=range(10), train_set=train_set, valid_set=valid_set
        )

    """
    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    training_set = ([inputs, targets],)
    validation_set = ([inputs, targets],)
    simple_brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )

    prof_simple = train(simple_brain, training_set, validation_set)
    # print(prof_simple.key_averages().table(sort_by="cpu_time_total"))
    # assert len(prof_simple.events()) >= 2500  # 2832 with torch==1.10.1
    # assert len(prof_simple.key_averages()) >= 60  # 72 with torch==1.10.1

    simple_brain_nitty_gritty = SimpleBrainNittyGritty(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    prof_nitty_gritty = train(
        simple_brain_nitty_gritty, training_set, validation_set
    )
    # print(prof_nitty_gritty.key_averages().table(sort_by="cpu_time_total"))
    # assert len(prof_nitty_gritty.events()) >= 2500  # 3030 with torch==1.10.1
    # assert len(prof_nitty_gritty.key_averages()) >= 60  # 74 with torch==1.10.1
    """

    # The outputs of this diff are only for visualisation, ``simple_delta._build_tree()`` will throw an error.
    """
    simple_delta, nitty_gritty_delta = events_diff(
        prof_simple.key_averages(), prof_nitty_gritty.key_averages()
    )
    # assert len(simple_delta) >= 4  # == 6  # before; config dependent: 7
    # assert len(nitty_gritty_delta) >= 4  # == 8  # before
    # assert simple_delta.total_average().count == 582 #Switching off becuase sometimes it fails
    # assert nitty_gritty_delta.total_average().count == 780 #Switching off becuase sometimes it fails
    with raises(Exception) as err_tree:
        simple_delta._build_tree()  # as mentioned.
    assert err_tree.type == AttributeError
    with raises(Exception) as err_averages:
        simple_delta.key_averages()  # as mentioned.
    assert err_averages.type == AssertionError
    " ""Both classes have alike numbers of function calls (given the same input data and train function).
    Sparing where both have the same number of calls:

    print(simple_delta.table(sort_by="cpu_time_total"))
    ----------------  ------------  ------------  ------------  ------------  ------------  ------------
                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    ----------------  ------------  ------------  ------------  ------------  ------------  ------------
         aten::zeros        17.44%     240.000us        27.69%     381.000us       6.350us            60
         aten::empty        26.96%     371.000us        27.11%     373.000us       1.492us           250
        aten::detach        18.68%     257.000us        25.65%     353.000us       5.694us            62
          aten::add_        25.00%     344.000us        25.00%     344.000us       5.931us            58
              detach         6.98%      96.000us         9.45%     130.000us       2.097us            62
         aten::zero_         4.94%      68.000us         4.94%      68.000us       0.756us            90
    ----------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 1.376ms

    print(nitty_gritty_delta.table(sort_by="cpu_time_total"))
    ----------------------  ------------  ------------  ------------  ------------  ------------  ------------
                      Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    ----------------------  ------------  ------------  ------------  ------------  ------------  ------------
        is this faster (?)        29.70%       1.024ms        76.19%       2.627ms     131.350us            20
               or that (?)        29.29%       1.010ms        67.00%       2.310ms     115.500us            20
               aten::zeros         8.32%     287.000us        15.52%     535.000us       5.350us           100
               aten::empty        14.01%     483.000us        14.07%     485.000us       1.470us           330
                aten::add_         9.92%     342.000us         9.92%     342.000us       5.700us            60
              aten::detach         2.81%      97.000us         6.09%     210.000us       3.500us            60
                    detach         3.28%     113.000us         4.00%     138.000us       2.300us            60
               aten::zero_         2.67%      92.000us         2.67%      92.000us       0.708us           130
    ----------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 3.448ms

    Curiosity doesn't come for free ;-)
    """


def test_scheduler(device):
    import torch
    from pytest import raises
    from torch.optim import SGD
    from speechbrain.core import Brain
    from speechbrain.utils.profiling import profile, schedule

    @schedule
    @profile
    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)
    test_set = ([inputs, targets],)

    # Profiling: __init__ constructor -- while scheduler: waiting --> nothing to report
    brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    assert brain.profiler.profiler is None
    assert len(brain.profiler.speechbrain_event_traces) == 0
    with raises(Exception) as err:
        brain.profiler.events()  # Tracing hasn't started, yet, so everything is in err. Scheduler says: wait.
    assert err.type == AssertionError
    assert brain.profiler.step_num == 0

    # Profiling: fit() for train operations.
    brain.fit(epoch_counter=range(10), train_set=train_set, valid_set=valid_set)
    assert brain.profiler.step_num == 20
    assert len(brain.profiler.speechbrain_event_traces) == 1
    # assert len(brain.profiler.events()) >= 250  # 293 with torch==1.10.1
    # assert len(brain.profiler.key_averages()) >= 60  # 73 with torch==1.10.1
    """print(brain.profiler.key_averages().table(sort_by="cpu_time_total"))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              ProfilerStep*        55.48%       1.504ms        99.00%       2.684ms       1.342ms             2
                                              aten::l1_loss         1.07%      29.000us         9.30%     252.000us      63.000us             4
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         4.57%     124.000us         7.01%     190.000us      47.500us             4
                                               aten::linear         2.32%      63.000us         6.93%     188.000us      94.000us             2
                                             aten::isfinite         0.89%      24.000us         5.35%     145.000us      48.333us             3
                                                   aten::to         1.00%      27.000us         4.46%     121.000us       7.562us            16
                                             aten::_to_copy         2.18%      59.000us         3.58%      97.000us       9.700us            10
                                                aten::stack         0.89%      24.000us         3.28%      89.000us      17.800us             5
                                     aten::l1_loss_backward         0.55%      15.000us         3.14%      85.000us      42.500us             2
                                                 aten::mean         0.48%      13.000us         3.10%      84.000us      42.000us             2
                                               aten::matmul         0.55%      15.000us         2.84%      77.000us      38.500us             2
                                                   aten::ne         0.77%      21.000us         2.73%      74.000us      24.667us             3
                                                   aten::mm         2.29%      62.000us         2.36%      64.000us      21.333us             3
                                                  aten::div         1.36%      37.000us         2.25%      61.000us      20.333us             3
                                    Optimizer.step#SGD.step         1.88%      51.000us         2.21%      60.000us      60.000us             1
    autograd::engine::evaluate_function: L1LossBackward0...         0.11%       3.000us         2.07%      56.000us      56.000us             1
                                                aten::zeros         1.44%      39.000us         1.99%      54.000us       6.750us             8
                                            L1LossBackward0         0.11%       3.000us         1.95%      53.000us      53.000us             1
                                                    aten::t         0.89%      24.000us         1.88%      51.000us      10.200us             5
                                                  aten::cat         0.59%      16.000us         1.81%      49.000us       9.800us             5
                                                 aten::div_         0.70%      19.000us         1.73%      47.000us      23.500us             2
           autograd::engine::evaluate_function: MmBackward0         0.11%       3.000us         1.70%      46.000us      46.000us             1
                                                  aten::abs         1.14%      31.000us         1.59%      43.000us       5.375us             8
                                                MmBackward0         0.37%      10.000us         1.59%      43.000us      43.000us             1
                                                  aten::sum         1.00%      27.000us         1.33%      36.000us      12.000us             3
                                                aten::empty         1.29%      35.000us         1.29%      35.000us       1.207us            29
          autograd::engine::evaluate_function: AddBackward0         0.63%      17.000us         1.25%      34.000us      34.000us             1
                                                 aten::_cat         0.85%      23.000us         1.22%      33.000us       6.600us             5
                                                 aten::add_         1.11%      30.000us         1.11%      30.000us       5.000us             6
                                            aten::transpose         0.70%      19.000us         0.92%      25.000us       5.000us             5
          autograd::engine::evaluate_function: DivBackward0         0.18%       5.000us         0.92%      25.000us      25.000us             1
                                                 aten::norm         0.77%      21.000us         0.89%      24.000us       8.000us             3
                          Optimizer.zero_grad#SGD.zero_grad         0.63%      17.000us         0.81%      22.000us      22.000us             1
                                                 aten::item         0.55%      15.000us         0.77%      21.000us       3.000us             7
                                                aten::copy_         0.77%      21.000us         0.77%      21.000us       2.100us            10
                                               DivBackward0         0.15%       4.000us         0.74%      20.000us      20.000us             1
    autograd::engine::evaluate_function: torch::autograd...         0.15%       4.000us         0.74%      20.000us      10.000us             2
                                                  aten::mul         0.55%      15.000us         0.74%      20.000us       5.000us             4
                                               aten::detach         0.37%      10.000us         0.70%      19.000us       3.167us             6
                                           aten::as_strided         0.59%      16.000us         0.63%      17.000us       1.062us            16
                                            aten::unsqueeze         0.41%      11.000us         0.59%      16.000us       2.667us             6
                                        aten::empty_strided         0.59%      16.000us         0.59%      16.000us       1.455us            11
                            torch::autograd::AccumulateGrad         0.30%       8.000us         0.59%      16.000us       8.000us             2
                                           aten::is_nonzero         0.18%       5.000us         0.59%      16.000us       5.333us             3
                                                 aten::view         0.55%      15.000us         0.55%      15.000us       3.000us             5
                                              aten::random_         0.52%      14.000us         0.52%      14.000us       7.000us             2
    autograd::engine::evaluate_function: UnsafeViewBackw...         0.11%       3.000us         0.52%      14.000us      14.000us             1
                                                  aten::sub         0.48%      13.000us         0.48%      13.000us       4.333us             3
                                                  aten::add         0.18%       5.000us         0.48%      13.000us      13.000us             1
                                            aten::ones_like         0.15%       4.000us         0.44%      12.000us      12.000us             1
                                                 aten::mul_         0.44%      12.000us         0.44%      12.000us       4.000us             3
                                                     detach         0.33%       9.000us         0.44%      12.000us       2.000us             6
                                        UnsafeViewBackward0         0.11%       3.000us         0.41%      11.000us      11.000us             1
                                                 aten::abs_         0.18%       5.000us         0.37%      10.000us       5.000us             2
                                           aten::empty_like         0.26%       7.000us         0.37%      10.000us       5.000us             2
                                                aten::clamp         0.26%       7.000us         0.37%      10.000us      10.000us             1
                                           aten::zeros_like         0.18%       5.000us         0.33%       9.000us       9.000us             1
                                                   aten::eq         0.33%       9.000us         0.33%       9.000us       3.000us             3
                                                aten::zero_         0.30%       8.000us         0.30%       8.000us       0.727us            11
                                         aten::_unsafe_view         0.22%       6.000us         0.30%       8.000us       4.000us             2
                                                aten::fill_         0.30%       8.000us         0.30%       8.000us       2.000us             4
                                              aten::reshape         0.18%       5.000us         0.30%       8.000us       8.000us             1
                                  aten::_local_scalar_dense         0.22%       6.000us         0.26%       7.000us       1.000us             7
            autograd::engine::evaluate_function: TBackward0         0.11%       3.000us         0.26%       7.000us       7.000us             1
                                              aten::resize_         0.22%       6.000us         0.22%       6.000us       1.200us             5
                                           aten::reciprocal         0.22%       6.000us         0.22%       6.000us       6.000us             1
                                                 aten::sgn_         0.15%       4.000us         0.15%       4.000us       4.000us             1
                                                 TBackward0         0.07%       2.000us         0.15%       4.000us       4.000us             1
                                       aten::_reshape_alias         0.11%       3.000us         0.11%       3.000us       3.000us             1
                                            aten::clamp_max         0.11%       3.000us         0.11%       3.000us       3.000us             1
                                         aten::resolve_conj         0.07%       2.000us         0.07%       2.000us       0.333us             6
                                    aten::broadcast_tensors         0.07%       2.000us         0.07%       2.000us       1.000us             2
                                               AddBackward0         0.04%       1.000us         0.04%       1.000us       1.000us             1
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 2.711ms  <===  above: Self CPU time total: 18.451ms (... the impact of warm-up)
    """

    @schedule
    @profile
    def train():
        # The step() function is executed inside speechbrain.core.brain.fit and is property of the Brain's profiler.
        # Above profiler and its scheduler are without power here, since prof.step() is not run - at all.
        brain.fit(
            epoch_counter=range(10), train_set=train_set, valid_set=valid_set
        )

    prof = train()
    # since we used the same brain (which has its own profiler)
    # assert brain.profiler.step_num == 20  # started again from 0 steps
    assert len(brain.profiler.speechbrain_event_traces) == 2
    # assert len(brain.profiler.events()) >= 250  # 293 with torch==1.10.1  # unchanged (overwritten with akin data)
    # assert len(brain.profiler.key_averages()) >= 60  # 73 with torch==1.10.1 # unchanged (akin data)
    # now, to the train function's profiler
    assert (
        prof.step_num == 0
    )  # the prof.step() operation wasn't run (not in scope) -> its scheduler is unawaken!
    assert not hasattr(prof, "speechbrain_event_traces")  # no trace collection
    with raises(Exception) as err_prof:
        prof.events()  # No tracing started with this one.
    assert err_prof.type == AssertionError  # sparing: key_averages()

    # But how to add profiling then if no writing access is there for a class... pretrained, for example:
    class SimpleBrainUntracked(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    brain_or_pretrained = SimpleBrainUntracked(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )

    # Set-up the profiler and hook it to the model.
    scheduled_profiler = schedule(profile)
    scheduled_profiler(brain_or_pretrained)

    # Profiling: still too early for scheduler!
    brain_or_pretrained.evaluate(test_set=test_set)  # -> step_num=1
    assert brain_or_pretrained.profiler.step_num == 1
    assert brain_or_pretrained.profiler.profiler is None

    # Profiling: scheduler warms-up.
    brain_or_pretrained.evaluate(
        test_set=(
            [inputs, targets],  # +1x test_set -> step_num=2
            [inputs, targets],
        )
    )
    assert brain_or_pretrained.profiler.step_num == 2
    # brain_or_pretrained.profiler.profiler will be set (not None anymore)
    # when run on cpu, there are no events - but cuda activities are recorded if existing
    # see: https://github.com/speechbrain/speechbrain/issues/1469
    if (
        torch.profiler.ProfilerActivity.CUDA
        in brain_or_pretrained.profiler.activities
    ):
        assert (
            len(
                set(
                    [
                        x.name
                        for x in brain_or_pretrained.profiler.profiler.function_events
                    ]
                )
                - {
                    "cudaGetDeviceCount",
                    "cudaGetDeviceProperties",
                    "cudaDeviceSynchronize",
                }
            )
        ) == 0
    else:
        assert len(brain_or_pretrained.profiler.events()) == 0

    # Profiling: scheduler warms-up...
    brain_or_pretrained.evaluate(
        test_set=(
            [inputs, targets],  # +1x test_set
            [inputs, targets],  # +2x test_set -> step_num=3
            [inputs, targets],
        )
    )
    assert brain_or_pretrained.profiler.step_num == 3
    if (
        torch.profiler.ProfilerActivity.CUDA
        in brain_or_pretrained.profiler.activities
    ):
        assert (
            len(
                set(
                    [
                        x.name
                        for x in brain_or_pretrained.profiler.profiler.function_events
                    ]
                )
                - {
                    "cudaGetDeviceCount",
                    "cudaGetDeviceProperties",
                    "cudaDeviceSynchronize",
                }
            )
        ) == 0
    else:
        assert len(brain_or_pretrained.profiler.events()) == 0

    # Profiling: first trace!
    brain_or_pretrained.evaluate(
        test_set=(
            [inputs, targets],  # +1x test_set
            [inputs, targets],  # +2x test_set
            [inputs, targets],  # +3x test_set -> step_num=4
            [inputs, targets],
        )
    )
    assert brain_or_pretrained.profiler.step_num == 4
    # assert len(brain_or_pretrained.profiler.events()) >= 4  # == 10  # before
    # assert len(brain_or_pretrained.profiler.key_averages()) >= 4  # == 5  # before
    assert (
        len(brain_or_pretrained.profiler.events()) >= 1
    )  # 1 on CPU; more w/ CUDA


def test_tracer(device):
    import torch
    from torch.optim import SGD
    from speechbrain.core import Brain
    from speechbrain.utils.profiling import profile, export

    @export
    @profile
    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)
    test_set = ([inputs, targets],)

    # Profiling: __init__ constructor and model training.
    brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    brain.fit(epoch_counter=range(10), train_set=train_set, valid_set=valid_set)

    # Pretrained example.
    class SimpleBrainUntracked(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    # No tracing during __init__
    brain_or_pretrained = SimpleBrainUntracked(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    profile(brain_or_pretrained, on_trace_ready=export(), with_stack=True)
    brain_or_pretrained.evaluate(test_set=test_set)

    # Set-up the profiler; hook it to the model, and benchmark inference.
    brain_or_pretrained2 = SimpleBrainUntracked(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    logged_profiler = export(profile)
    assert brain_or_pretrained2.profiler is None
    logged_profiler(brain_or_pretrained2)
    brain_or_pretrained2.evaluate(test_set=test_set)


def test_aggregated_traces(device):
    import torch
    from torch.optim import SGD
    from speechbrain.core import Brain
    from speechbrain.utils.profiling import profile

    @profile
    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)
    test_set = (
        [inputs, targets],
        [inputs, targets],
    )

    # Profiling: __init__ constructor -- while scheduler: waiting --> nothing to report
    brain = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )

    # Profiling: empty traces
    assert len(brain.profiler.speechbrain_event_traces) == 1
    """
    init_report = brain.profiler.merge_traces()
    assert len(init_report) >= 1
    # assert len(init_report) >= 4  # == 6  # before; config dependent: 7
    assert len(brain.profiler.speechbrain_event_traces) == 1
    " ""print(brain.profiler.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                           Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
        aten::_has_compatible_shallow_copy_type        80.00%      12.000us        80.00%      12.000us       3.000us             4
                                       aten::to        20.00%       3.000us        20.00%       3.000us       1.500us             2
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 15.000us

    print(init_report.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                           Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
        aten::_has_compatible_shallow_copy_type        80.00%      12.000us        80.00%      12.000us       3.000us             4
                                       aten::to        20.00%       3.000us        20.00%       3.000us       1.500us             2
    -------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 15.000us
    """

    brain.fit(epoch_counter=range(10), train_set=train_set, valid_set=valid_set)
    assert len(brain.profiler.speechbrain_event_traces) == 2
    # assert len(brain.profiler.speechbrain_event_traces[0]) >= 4  # == 6  # before; config dependent: 7
    # assert len(brain.profiler.speechbrain_event_traces[1]) >= 2500  # 2862 with torch==1.10.1
    # assert len(brain.profiler.events()) >= 2500  # 2832 with torch==1.10.1
    # assert len(brain.profiler.events().key_averages()) >= 60  # 72 with torch==1.10.1
    """print(brain.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              aten::l1_loss         2.39%     415.000us        25.28%       4.392ms     109.800us            40
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        12.65%       2.198ms        20.06%       3.485ms      87.125us            40
                                               aten::linear         1.41%     245.000us        13.24%       2.299ms     114.950us            20
                                                   aten::to         2.04%     354.000us        10.59%       1.839ms      11.494us           160
                                             aten::isfinite         2.12%     369.000us         9.87%       1.714ms      57.133us            30
                                                 aten::mean         1.13%     196.000us         9.15%       1.589ms      79.450us            20
                                                aten::stack         2.33%     404.000us         8.94%       1.553ms      31.060us            50
                                             aten::_to_copy         5.67%     985.000us         8.67%       1.506ms      15.060us           100
                                               aten::matmul         1.57%     273.000us         7.83%       1.360ms      68.000us            20
                                     aten::l1_loss_backward         1.22%     212.000us         6.67%       1.158ms      57.900us            20
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 17.370ms
    """

    # Profiling: aggregate traces
    """
    short_report = brain.profiler.merge_traces()
    assert len(short_report) >= 1
    # assert len(short_report) >= 2500  # 2838 with torch==1.10.1
    # assert len(short_report.key_averages()) >= 60  # 73 with torch==1.10.1
    " ""print(short_report.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              aten::l1_loss         2.39%     415.000us        25.26%       4.392ms     109.800us            40
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        12.64%       2.198ms        20.05%       3.485ms      87.125us            40
                                               aten::linear         1.41%     245.000us        13.22%       2.299ms     114.950us            20
                                                   aten::to         2.05%     357.000us        10.60%       1.842ms      11.370us           162
                                             aten::isfinite         2.12%     369.000us         9.86%       1.714ms      57.133us            30
                                                 aten::mean         1.13%     196.000us         9.14%       1.589ms      79.450us            20
                                                aten::stack         2.32%     404.000us         8.93%       1.553ms      31.060us            50
                                             aten::_to_copy         5.67%     985.000us         8.66%       1.506ms      15.060us           100
                                               aten::matmul         1.57%     273.000us         7.82%       1.360ms      68.000us            20
                                     aten::l1_loss_backward         1.22%     212.000us         6.66%       1.158ms      57.900us            20
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 17.385ms
    """

    brain.evaluate(test_set=test_set)
    brain.evaluate(test_set=test_set)
    brain.evaluate(test_set=test_set)
    assert len(brain.profiler.speechbrain_event_traces) == 5
    # assert len(brain.profiler.speechbrain_event_traces[0]) >= 4  # == 6  # before; config dependent: 7
    # assert len(brain.profiler.speechbrain_event_traces[1]) >= 2500  # 2862 with torch==1.10.1
    # assert len(brain.profiler.speechbrain_event_traces[2]) >= 125  # 143 with torch==1.10.1
    # assert len(brain.profiler.speechbrain_event_traces[3]) >= 125  # 143 with torch==1.10.1
    # assert len(brain.profiler.speechbrain_event_traces[4]) >= 125  # 143 with torch==1.10.1
    # assert len(brain.profiler.events()) >= 125  # 141 with torch==1.10.1
    # assert len(brain.profiler.events().key_averages()) >= 25  # 42 with torch==1.10.1
    # the following is only for the last call of the 3x brain.evaluate()
    """print(brain.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              aten::l1_loss         3.54%      23.000us        37.38%     243.000us      60.750us             4
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        16.62%     108.000us        29.38%     191.000us      63.667us             3
                                               aten::linear         2.77%      18.000us        20.46%     133.000us      66.500us             2
                                             aten::isfinite         3.54%      23.000us        14.31%      93.000us      46.500us             2
                                                 aten::mean         1.85%      12.000us        12.62%      82.000us      41.000us             2
                                                aten::stack         3.23%      21.000us        12.15%      79.000us      19.750us             4
                                               aten::matmul         2.62%      17.000us        11.69%      76.000us      38.000us             2
                                                 aten::div_         2.92%      19.000us         7.54%      49.000us      24.500us             2
                                                   aten::to         1.85%      12.000us         7.08%      46.000us       7.667us             6
                                                   aten::mm         6.31%      41.000us         6.77%      44.000us      22.000us             2
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 650.000us
    """

    # Profiling: putting previous benchmark reporting together.
    """
    full_report = brain.profiler.merge_traces()
    assert len(full_report) >= 1
    # assert len(full_report.key_averages()) >= 60  # 73 with torch==1.10.1
    # In this minimal example, only 73 functions matter.
    # Some events are duplicated (perhaps from wrapping functions):
    # => they appear stacked & EventList._remove_dup_nodes drops direct child events of same name as their parent.
    num_events = sum([len(x) for x in brain.profiler.speechbrain_event_traces])
    assert num_events >= 1
    # assert num_events >= 3000  # 3297 with torch==1.10.1  # expected: 6 + 2862 + 3x143 = 3297
    # Apparently, this depends on how this test is run (by its own or as part of the entire file's test suite).
    # assert (num_events == len(full_report)) or (len(full_report) == len(set([x.id for x in full_report])))
    # ... not tested, why
    " ""print(full_report.key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              aten::l1_loss         2.21%     427.000us        27.58%       5.326ms     102.423us            52
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        11.89%       2.297ms        21.80%       4.210ms      85.918us            49
                                               aten::linear         1.57%     304.000us        14.61%       2.821ms     108.500us            26
                                             aten::isfinite         2.53%     488.000us        10.72%       2.071ms      57.528us            36
                                                   aten::to         2.03%     392.000us        10.42%       2.013ms      11.183us           180
                                                 aten::mean         1.15%     223.000us         9.81%       1.894ms      72.846us            26
                                                aten::stack         2.34%     452.000us         9.54%       1.842ms      29.710us            62
                                             aten::_to_copy         5.49%       1.061ms         8.51%       1.643ms      14.670us           112
                                               aten::matmul         1.65%     318.000us         8.48%       1.638ms      63.000us            26
                                                 aten::div_         1.50%     290.000us         6.95%       1.343ms      47.964us            28
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 19.311ms
    """
    # 19.311ms = 3x ~650.000us + 17.370ms + 15.000us <=> 1st & 2nd call of brain.evaluate() = 1276us = 2x 638us
    # max([x.time_range.end for x in full_report]) -> 41965 (us)


def test_profile_details(device):
    import torch

    # from copy import deepcopy
    from torch.optim import SGD
    from speechbrain.core import Brain
    from speechbrain.utils.profiling import (
        profile_analyst,
        profile_optimiser,
        export,
        # events_diff,
    )

    class SimpleBrain(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    model = torch.nn.Linear(in_features=10, out_features=10, device=device)
    inputs = torch.rand(10, 10, device=device)
    targets = torch.rand(10, 10, device=device)
    train_set = ([inputs, targets],)
    valid_set = ([inputs, targets],)
    test_set = (
        [inputs, targets],
        [inputs, targets],
        [inputs, targets],
        [inputs, targets],
        [inputs, targets],
        [inputs, targets],
    )

    brain_analyst = profile_analyst(
        SimpleBrain(
            {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
        )
    )
    brain_optimiser = profile_optimiser(
        SimpleBrain(
            {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
        )
    )

    assert len(brain_analyst.profiler.speechbrain_event_traces) == 0
    brain_analyst.fit(
        epoch_counter=range(10), train_set=train_set, valid_set=valid_set
    )
    assert len(brain_analyst.profiler.speechbrain_event_traces) == 1
    # assert len(brain_analyst.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
    # assert len(brain_analyst.profiler.events()) >= 250  # 293 with torch==1.10.1
    # assert len(brain_analyst.profiler.events().key_averages()) >= 60  # 73 with torch==1.10.1
    """print(brain_analyst.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls   Total FLOPs
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              ProfilerStep*        36.62%       4.345ms        98.50%      11.686ms       5.843ms         796 b      -1.66 Kb             2            --
                                              aten::l1_loss         2.50%     297.000us        14.37%       1.705ms     426.250us          16 b        -800 b             4            --
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         3.91%     464.000us        11.01%       1.306ms     326.500us       1.55 Kb         -80 b             4            --
                                             aten::isfinite         3.10%     368.000us         9.65%       1.145ms     381.667us           3 b         -18 b             3            --
                                                aten::stack         2.79%     331.000us         8.46%       1.004ms     200.800us       1.57 Kb           0 b             5            --
                                                   aten::to         1.96%     232.000us         7.25%     860.000us      53.750us          40 b           0 b            16            --
                                               aten::linear         1.40%     166.000us         6.55%     777.000us     388.500us         800 b           0 b             2            --
                                     aten::l1_loss_backward         1.58%     188.000us         6.08%     721.000us     360.500us         400 b          -4 b             2            --
                                             aten::_to_copy         4.28%     508.000us         5.29%     628.000us      62.800us          40 b           0 b            10            --
                                                 aten::mean         0.89%     105.000us         4.44%     527.000us     263.500us           8 b           8 b             2            --
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 11.864ms
    """
    # 6-batch inference
    brain_analyst.evaluate(test_set=test_set)
    assert len(brain_analyst.profiler.speechbrain_event_traces) == 2
    # assert len(brain_analyst.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
    # assert len(brain_analyst.profiler.speechbrain_event_traces[1]) >= 125  # 144 with torch==1.10.1
    # as of evaluate() call
    # assert len(brain_analyst.profiler.events()) >= 125  # 142 with torch==1.10.1
    # assert len(brain_analyst.profiler.events().key_averages()) >= 25  # 42 with torch==1.10.1
    """print(brain_analyst.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls   Total FLOPs
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              ProfilerStep*        19.24%       1.129ms        96.92%       5.687ms       2.844ms         796 b      -1.61 Kb             2            --
                                              aten::l1_loss         5.16%     303.000us        35.50%       2.083ms     520.750us          16 b        -800 b             4            --
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         7.41%     435.000us        25.95%       1.523ms     761.500us       1.55 Kb         -40 b             2            --
                                                aten::stack         6.15%     361.000us        18.20%       1.068ms     267.000us       1.56 Kb           0 b             4            --
                                               aten::linear         3.43%     201.000us        15.78%     926.000us     463.000us         800 b           0 b             2            --
                                                 aten::mean         2.42%     142.000us        11.93%     700.000us     350.000us           8 b           8 b             2            --
                                             aten::isfinite         3.72%     218.000us        10.84%     636.000us     318.000us           2 b         -12 b             2            --
                                                  aten::cat         2.68%     157.000us         8.95%     525.000us     131.250us       1.56 Kb           0 b             4            --
                                               aten::matmul         3.83%     225.000us         8.88%     521.000us     260.500us         800 b           0 b             2            --
                                                 aten::div_         3.34%     196.000us         6.97%     409.000us     204.500us           0 b          -8 b             2            --
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 5.868ms
    """

    brain_optimiser.fit(
        epoch_counter=range(10), train_set=train_set, valid_set=valid_set
    )
    # key_avg_fit = deepcopy(brain_optimiser.profiler.events().key_averages())
    """print(brain_optimiser.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              ProfilerStep*        50.73%       1.874ms        98.86%       3.652ms       1.826ms         796 b      -1.66 Kb             2
                                              aten::l1_loss         1.38%      51.000us        11.83%     437.000us     109.250us          16 b        -400 b             4
                                             aten::isfinite         1.49%      55.000us         8.69%     321.000us     107.000us           3 b         -16 b             3
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         4.03%     149.000us         7.58%     280.000us      70.000us       1.55 Kb         -64 b             4
                                               aten::linear         0.51%      19.000us         5.66%     209.000us     104.500us         800 b           0 b             2
                                                  aten::abs         3.19%     118.000us         5.58%     206.000us      25.750us          24 b          12 b             8
                                     aten::l1_loss_backward         0.92%      34.000us         4.28%     158.000us      79.000us         400 b          -4 b             2
                                                aten::stack         1.00%      37.000us         3.87%     143.000us      28.600us       1.57 Kb           0 b             5
                                                aten::empty         3.76%     139.000us         3.76%     139.000us       4.793us         544 b         544 b            29
                                                   aten::to         0.76%      28.000us         3.76%     139.000us       8.688us          44 b           4 b            16
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 3.694ms
    # to 11.864ms (analyst)
    """

    brain_optimiser.evaluate(test_set=test_set)
    """
    key_avg_evaluate = deepcopy(
        brain_optimiser.profiler.events().key_averages()
    )
    """
    """print(brain_optimiser.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                              ProfilerStep*        24.80%     524.000us        96.50%       2.039ms       1.020ms         796 b      -1.61 Kb             2
                                              aten::l1_loss         2.74%      58.000us        33.65%     711.000us     177.750us          16 b        -800 b             4
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         9.94%     210.000us        21.11%     446.000us     223.000us       1.55 Kb         -40 b             2
                                             aten::isfinite         3.64%      77.000us        15.76%     333.000us     166.500us           2 b         -12 b             2
                                               aten::linear         1.04%      22.000us        11.88%     251.000us     125.500us         800 b           0 b             2
                                                 aten::mean         2.04%      43.000us        11.74%     248.000us     124.000us           8 b           8 b             2
                                                aten::stack         3.31%      70.000us        10.18%     215.000us      53.750us       1.56 Kb           0 b             4
                                                 aten::div_         4.83%     102.000us         7.90%     167.000us      83.500us           0 b          -8 b             2
                                               aten::matmul         1.61%      34.000us         7.38%     156.000us      78.000us         800 b           0 b             2
                                                   aten::ne         4.54%      96.000us         6.72%     142.000us      71.000us           2 b          -6 b             2
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 2.113ms
    # to 5.868ms (analyst)
    """
    # same check as for analyst
    assert len(brain_optimiser.profiler.speechbrain_event_traces) == 2
    # assert len(brain_optimiser.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
    # assert len(brain_optimiser.profiler.speechbrain_event_traces[1]) >= 125  # 144 with torch==1.10.1
    # as of evaluate() call
    # assert len(brain_optimiser.profiler.events()) >= 125  # 142 with torch==1.10.1
    # assert len(brain_optimiser.profiler.events().key_averages()) >= 25  # 42 with torch==1.10.1
    # different config
    assert (
        brain_optimiser.profiler.record_shapes
        != brain_analyst.profiler.record_shapes
    )
    assert (
        brain_optimiser.profiler.with_stack != brain_analyst.profiler.with_stack
    )
    assert (
        brain_optimiser.profiler.with_flops != brain_analyst.profiler.with_flops
    )
    # same config
    assert (
        brain_optimiser.profiler.with_modules
        == brain_analyst.profiler.with_modules
    )
    assert (
        brain_optimiser.profiler.profile_memory
        == brain_analyst.profiler.profile_memory
    )

    """
    # let's take a look at the diff
    diff_fit, diff_evaluate = events_diff(key_avg_fit, key_avg_evaluate)
    # assert len(diff_fit) >= 50  # 64 with torch==1.10.1
    # assert len(diff_evaluate) >= 25  # 33 with torch==1.10.1
    # assert diff_fit.total_average().count >= 250  # 273 with torch==1.10.1
    # assert diff_evaluate.total_average().count >= 100  # 122 with torch==1.10.1
    " ""For curiosity only... the printed FunctionEvents differ by (name, # of Calls)
    print(diff_fit.table(sort_by="cpu_time_total"))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                             aten::isfinite         3.35%      55.000us        19.55%     321.000us     107.000us           3 b         -16 b             3
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         9.07%     149.000us        17.05%     280.000us      70.000us       1.55 Kb         -64 b             4
                                                  aten::abs         7.19%     118.000us        12.55%     206.000us      25.750us          24 b          12 b             8
                                     aten::l1_loss_backward         2.07%      34.000us         9.62%     158.000us      79.000us         400 b          -4 b             2
                                                aten::stack         2.25%      37.000us         8.71%     143.000us      28.600us       1.57 Kb           0 b             5
                                                aten::empty         8.47%     139.000us         8.47%     139.000us       4.793us         544 b         544 b            29
                                                   aten::to         1.71%      28.000us         8.47%     139.000us       8.688us          44 b           4 b            16
                                             aten::_to_copy         1.89%      31.000us         7.00%     115.000us      11.500us          40 b           0 b            10
                                                   aten::mm         6.46%     106.000us         6.58%     108.000us      36.000us       1.17 Kb       1.17 Kb             3
                                                   aten::ne         4.26%      70.000us         6.46%     106.000us      35.333us           3 b          -9 b             3
                                                aten::zeros         1.95%      32.000us         6.03%      99.000us      12.375us          32 b           0 b             8
    autograd::engine::evaluate_function: L1LossBackward0...         0.37%       6.000us         5.97%      98.000us      98.000us         396 b          -4 b             1
                                            L1LossBackward0         0.24%       4.000us         5.60%      92.000us      92.000us         400 b           0 b             1
                                                  aten::cat         0.97%      16.000us         5.05%      83.000us      16.600us       1.57 Kb           0 b             5
                                                  aten::div         2.80%      46.000us         4.81%      79.000us      26.333us          12 b           0 b             3
                                    Optimizer.step#SGD.step         3.71%      61.000us         4.51%      74.000us      74.000us          -4 b         -20 b             1
                                                 aten::_cat         2.13%      35.000us         4.08%      67.000us      13.400us       1.57 Kb           0 b             5
                                        aten::empty_strided         3.41%      56.000us         3.41%      56.000us       5.091us          44 b          44 b            11
           autograd::engine::evaluate_function: MmBackward0         0.37%       6.000us         3.41%      56.000us      56.000us           0 b        -400 b             1
                                                    aten::t         2.13%      35.000us         3.11%      51.000us      10.200us           0 b           0 b             5
                                                 aten::add_         3.05%      50.000us         3.05%      50.000us       8.333us           0 b           0 b             6
                                                MmBackward0         0.43%       7.000us         3.05%      50.000us      50.000us         400 b           0 b             1
                                                  aten::mul         2.01%      33.000us         2.62%      43.000us      10.750us           5 b          -3 b             4
                                                  aten::sum         2.01%      33.000us         2.50%      41.000us      13.667us          40 b          40 b             3
          autograd::engine::evaluate_function: DivBackward0         0.55%       9.000us         2.38%      39.000us      39.000us          -4 b          -8 b             1
                                                 aten::norm         2.13%      35.000us         2.38%      39.000us      13.000us          12 b          12 b             3
                          Optimizer.zero_grad#SGD.zero_grad         1.46%      24.000us         1.89%      31.000us      31.000us          -4 b          -4 b             1
                                               DivBackward0         0.24%       4.000us         1.83%      30.000us      30.000us           4 b           0 b             1
                                                 aten::item         1.04%      17.000us         1.77%      29.000us       4.143us           0 b           0 b             7
          autograd::engine::evaluate_function: AddBackward0         0.43%       7.000us         1.77%      29.000us      29.000us          40 b           0 b             1
                                              aten::random_         1.71%      28.000us         1.71%      28.000us      14.000us           0 b           0 b             2
                                                  aten::sub         1.71%      28.000us         1.71%      28.000us       9.333us         800 b         800 b             3
                                                aten::copy_         1.71%      28.000us         1.71%      28.000us       2.800us           0 b           0 b            10
    autograd::engine::evaluate_function: torch::autograd...         0.30%       5.000us         1.71%      28.000us      14.000us        -440 b           0 b             2
                                              aten::resize_         1.64%      27.000us         1.64%      27.000us       5.400us       1.57 Kb       1.57 Kb             5
                                                  aten::add         0.85%      14.000us         1.64%      27.000us      27.000us           4 b           0 b             1
                                                   aten::eq         1.58%      26.000us         1.58%      26.000us       8.667us           3 b           3 b             3
                                            aten::unsqueeze         0.91%      15.000us         1.40%      23.000us       3.833us           0 b           0 b             6
                            torch::autograd::AccumulateGrad         0.97%      16.000us         1.40%      23.000us      11.500us        -440 b        -440 b             2
                                           aten::is_nonzero         0.43%       7.000us         1.40%      23.000us       7.667us           0 b           0 b             3
                                               aten::detach         0.55%       9.000us         1.28%      21.000us       3.500us           0 b           0 b             6
                                           aten::as_strided         1.16%      19.000us         1.16%      19.000us       1.188us           0 b           0 b            16
                                                 aten::view         1.04%      17.000us         1.04%      17.000us       3.400us           0 b           0 b             5
                                            aten::transpose         0.67%      11.000us         0.97%      16.000us       3.200us           0 b           0 b             5
                                           aten::empty_like         0.37%       6.000us         0.91%      15.000us       7.500us         404 b           0 b             2
                                                     detach         0.73%      12.000us         0.85%      14.000us       2.333us           0 b           0 b             6
                                                aten::clamp         0.61%      10.000us         0.85%      14.000us      14.000us           4 b           4 b             1
                                  aten::_local_scalar_dense         0.73%      12.000us         0.79%      13.000us       1.857us           0 b           0 b             7
                                                 aten::mul_         0.79%      13.000us         0.79%      13.000us       4.333us           0 b           0 b             3
                                            aten::ones_like         0.24%       4.000us         0.73%      12.000us      12.000us           4 b           0 b             1
                                           aten::zeros_like         0.30%       5.000us         0.73%      12.000us      12.000us         400 b           0 b             1
    autograd::engine::evaluate_function: UnsafeViewBackw...         0.18%       3.000us         0.73%      12.000us      12.000us           0 b           0 b             1
                                        UnsafeViewBackward0         0.12%       2.000us         0.55%       9.000us       9.000us           0 b           0 b             1
                                           aten::reciprocal         0.49%       8.000us         0.49%       8.000us       8.000us           4 b           4 b             1
                                              aten::reshape         0.24%       4.000us         0.43%       7.000us       7.000us           0 b           0 b             1
            autograd::engine::evaluate_function: TBackward0         0.12%       2.000us         0.43%       7.000us       7.000us           0 b           0 b             1
                                                aten::zero_         0.37%       6.000us         0.37%       6.000us       0.545us           0 b           0 b            11
                                                aten::fill_         0.37%       6.000us         0.37%       6.000us       1.500us           0 b           0 b             4
                                                 aten::sgn_         0.30%       5.000us         0.30%       5.000us       5.000us           0 b           0 b             1
                                                 TBackward0         0.06%       1.000us         0.30%       5.000us       5.000us           0 b           0 b             1
                                            aten::clamp_max         0.24%       4.000us         0.24%       4.000us       4.000us           0 b           0 b             1
                                       aten::_reshape_alias         0.18%       3.000us         0.18%       3.000us       3.000us           0 b           0 b             1
                                         aten::resolve_conj         0.12%       2.000us         0.12%       2.000us       0.333us           0 b           0 b             6
                                               AddBackward0         0.06%       1.000us         0.06%       1.000us       1.000us           0 b           0 b             1
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 1.642ms


    print(diff_evaluate.table(sort_by="cpu_time_total"))
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        16.03%     210.000us        34.05%     446.000us     223.000us       1.55 Kb         -40 b             2
                                             aten::isfinite         5.88%      77.000us        25.42%     333.000us     166.500us           2 b         -12 b             2
                                                aten::stack         5.34%      70.000us        16.41%     215.000us      53.750us       1.56 Kb           0 b             4
                                                   aten::ne         7.33%      96.000us        10.84%     142.000us      71.000us           2 b          -6 b             2
                                                aten::empty        10.08%     132.000us        10.08%     132.000us       8.250us          80 b          80 b            16
                                                   aten::to         1.37%      18.000us         8.63%     113.000us      18.833us          16 b           0 b             6
                                                aten::zeros         3.13%      41.000us         8.55%     112.000us      28.000us          16 b           0 b             4
                                                  aten::cat         1.45%      19.000us         8.40%     110.000us      27.500us       1.56 Kb           0 b             4
                                                   aten::mm         7.02%      92.000us         7.33%      96.000us      48.000us         800 b         800 b             2
                                                  aten::abs         4.58%      60.000us         7.33%      96.000us      16.000us          16 b           8 b             6
                                             aten::_to_copy         2.21%      29.000us         7.25%      95.000us      23.750us          16 b           0 b             4
                                                 aten::_cat         2.82%      37.000us         6.95%      91.000us      22.750us       1.56 Kb           0 b             4
                                              aten::resize_         3.51%      46.000us         3.51%      46.000us      11.500us       1.56 Kb       1.56 Kb             4
                                        aten::empty_strided         3.36%      44.000us         3.36%      44.000us      11.000us          16 b          16 b             4
                                                  aten::sub         2.98%      39.000us         2.98%      39.000us      19.500us         800 b         800 b             2
                                                  aten::sum         2.14%      28.000us         2.90%      38.000us      19.000us           0 b           0 b             2
                                                 aten::add_         2.82%      37.000us         2.82%      37.000us      18.500us           0 b           0 b             2
                                                    aten::t         1.68%      22.000us         2.75%      36.000us      18.000us           0 b           0 b             2
                                            aten::unsqueeze         1.98%      26.000us         2.67%      35.000us       8.750us           0 b           0 b             4
                                                   aten::eq         2.37%      31.000us         2.37%      31.000us      15.500us           2 b           2 b             2
                                                  aten::mul         2.37%      31.000us         2.37%      31.000us      15.500us           2 b           2 b             2
                                           aten::is_nonzero         0.53%       7.000us         1.76%      23.000us      11.500us           0 b           0 b             2
                                                 aten::item         0.92%      12.000us         1.76%      23.000us       5.750us           0 b           0 b             4
                                                aten::copy_         1.68%      22.000us         1.68%      22.000us       5.500us           0 b           0 b             4
                                           aten::as_strided         1.37%      18.000us         1.37%      18.000us       2.250us           0 b           0 b             8
                                                 aten::view         1.30%      17.000us         1.30%      17.000us       4.250us           0 b           0 b             4
                                               aten::detach         0.31%       4.000us         1.15%      15.000us       7.500us           0 b           0 b             2
                                            aten::transpose         0.69%       9.000us         1.07%      14.000us       7.000us           0 b           0 b             2
                                                     detach         0.84%      11.000us         0.84%      11.000us       5.500us           0 b           0 b             2
                                  aten::_local_scalar_dense         0.84%      11.000us         0.84%      11.000us       2.750us           0 b           0 b             4
                                                aten::fill_         0.46%       6.000us         0.46%       6.000us       3.000us           0 b           0 b             2
                                                aten::zero_         0.31%       4.000us         0.31%       4.000us       1.000us           0 b           0 b             4
                                         aten::resolve_conj         0.31%       4.000us         0.31%       4.000us       1.000us           0 b           0 b             4
    -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
    Self CPU time total: 1.310ms
    """

    # set hook afterwards
    brain_analyst_raw = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    brain_optimiser_raw = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    brain_analyst_raw.fit(
        epoch_counter=range(10), train_set=train_set, valid_set=valid_set
    )
    profile_analyst(brain_analyst_raw)
    brain_analyst_raw.evaluate(test_set=test_set)
    assert getattr(brain_analyst_raw.profiler, "record_shapes") is True
    assert getattr(brain_analyst_raw.profiler, "with_stack") is True
    assert getattr(brain_analyst_raw.profiler, "with_flops") is True

    brain_optimiser_raw.fit(
        epoch_counter=range(10), train_set=train_set, valid_set=valid_set
    )
    profile_optimiser(brain_optimiser_raw)
    brain_optimiser_raw.evaluate(test_set=test_set)
    assert getattr(brain_optimiser_raw.profiler, "record_shapes") is False
    assert getattr(brain_optimiser_raw.profiler, "with_stack") is False
    assert getattr(brain_optimiser_raw.profiler, "with_flops") is False

    # wrap functions
    @profile_analyst
    def train_analyst(brain: SimpleBrain):
        brain.fit(
            epoch_counter=range(10), train_set=train_set, valid_set=valid_set
        )

    @export
    @profile_optimiser
    def evaluate_optimiser(brain: SimpleBrain):
        brain.evaluate(test_set=test_set)

    brain_raw = SimpleBrain(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    assert brain_raw.profiler is None
    train_analyst(brain_raw)
    assert brain_raw.profiler is None
    evaluate_optimiser(brain_raw)
    assert brain_raw.profiler is None

    # profile classes
    @export
    @profile_analyst
    class SimpleBrainAnalyst(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    @profile_optimiser
    class SimpleBrainOptimiser(Brain):
        def compute_forward(self, batch, stage):
            return self.modules.model(batch[0])

        def compute_objectives(self, predictions, batch, stage):
            return torch.nn.functional.l1_loss(predictions, batch[1])

    simple_brain_analyst = SimpleBrainAnalyst(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    assert getattr(simple_brain_analyst.profiler, "record_shapes") is True
    assert getattr(simple_brain_analyst.profiler, "with_stack") is True
    assert getattr(simple_brain_analyst.profiler, "with_flops") is True
    simple_brain_analyst.evaluate(test_set=test_set)

    simple_brain_optimiser = SimpleBrainOptimiser(
        {"model": model}, lambda x: SGD(x, 0.1), run_opts={"device": device}
    )
    assert getattr(simple_brain_optimiser.profiler, "record_shapes") is False
    assert getattr(simple_brain_optimiser.profiler, "with_stack") is False
    assert getattr(simple_brain_optimiser.profiler, "with_flops") is False
    simple_brain_optimiser.fit(
        epoch_counter=range(10), train_set=train_set, valid_set=valid_set
    )

tests/unittests/test_profiling.py
TzuchengChang/NASS
import torch


def test_ConcatDatasetBatchSampler(device):
    from torch.utils.data import TensorDataset, ConcatDataset, DataLoader
    from speechbrain.dataio.sampler import (
        ReproducibleRandomSampler,
        ConcatDatasetBatchSampler,
    )
    import numpy as np

    datasets = []
    for i in range(3):
        if i == 0:
            datasets.append(
                TensorDataset(torch.arange(i * 10, (i + 1) * 10, device=device))
            )
        else:
            datasets.append(
                TensorDataset(torch.arange(i * 6, (i + 1) * 6, device=device))
            )

    samplers = [ReproducibleRandomSampler(x) for x in datasets]
    dataset = ConcatDataset(datasets)
    loader = DataLoader(
        dataset, batch_sampler=ConcatDatasetBatchSampler(samplers, [1, 1, 1])
    )

    concat_data = []

    for data in loader:
        concat_data.append([x.item() for x in data[0]])
    concat_data = np.array(concat_data)

    non_cat_data = []
    for i in range(len(samplers)):
        c_data = []
        loader = DataLoader(dataset.datasets[i], sampler=samplers[i])

        for data in loader:
            c_data.append(data[0].item())

        non_cat_data.append(c_data)

    minlen = min([len(x) for x in non_cat_data])
    non_cat_data = [x[:minlen] for x in non_cat_data]
    non_cat_data = np.array(non_cat_data)
    np.testing.assert_array_equal(non_cat_data.T, concat_data)

tests/unittests/test_samplers.py
TzuchengChang/NASS
import torch


def test_normalize(device):

    from speechbrain.processing.signal_processing import compute_amplitude
    from speechbrain.processing.signal_processing import rescale
    import random
    import numpy as np

    for scale in ["dB", "linear"]:
        for amp_type in ["peak", "avg"]:
            for test_vec in [
                torch.zeros((100), device=device),
                torch.rand((10, 100), device=device),
                torch.rand((10, 100, 5), device=device),
            ]:

                lengths = (
                    test_vec.size(1)
                    if len(test_vec.shape) > 1
                    else test_vec.size(0)
                )
                amp = compute_amplitude(test_vec, lengths, amp_type, scale)
                scaled_back = rescale(
                    random.random() * test_vec, lengths, amp, amp_type, scale
                )
                np.testing.assert_array_almost_equal(
                    scaled_back.cpu().numpy(), test_vec.cpu().numpy()
                )

tests/unittests/test_signal_processing.py
TzuchengChang/NASS
import os
import torch


def test_tokenizer():
    from speechbrain.tokenizers.SentencePiece import SentencePiece

    gt = [
        ["HELLO", "MORNING", "MORNING", "HELLO"],
        ["HELLO", "MORNING", "HELLO"],
    ]

    # Word-level input test
    dict_int2lab = {1: "HELLO", 2: "MORNING"}

    spm = SentencePiece(
        os.path.abspath("tests/tmp/tokenizer_data"),
        100,
        annotation_train=os.path.abspath(
            "tests/samples/annotation/tokenizer.csv"
        ),
        annotation_read="wrd",
        model_type="bpe",
    )
    encoded_seq_ids, encoded_seq_pieces = spm(
        torch.Tensor([[1, 2, 2, 1], [1, 2, 1, 0]]),
        torch.Tensor([1.0, 0.75]),
        dict_int2lab,
        task="encode",
    )
    lens = (encoded_seq_pieces * encoded_seq_ids.shape[1]).round().int()
    # decode from torch tensors (batch, batch_lens)
    words_seq = spm(encoded_seq_ids, encoded_seq_pieces, task="decode")
    assert words_seq == gt, "output not the same"
    # decode from a list of bpe sequence (without padding)
    hyps_list = [
        encoded_seq_ids[0].int().tolist(),
        encoded_seq_ids[1][: lens[1]].int().tolist(),
    ]
    words_seq = spm(hyps_list, task="decode_from_list")
    assert words_seq == gt, "output not the same"

    # Char-level input test
    dict_int2lab = {
        1: "H",
        2: "E",
        3: "L",
        4: "O",
        5: "M",
        6: "R",
        7: "N",
        8: "I",
        9: "G",
        10: "_",
    }

    spm = SentencePiece(
        os.path.abspath("tests/tmp/tokenizer_data"),
        100,
        annotation_train=os.path.abspath(
            "tests/sample/annotation/tokenzer.csv"
        ),
        annotation_read="char",
        char_format_input=True,
        model_type="bpe",
    )
    encoded_seq_ids, encoded_seq_pieces = spm(
        torch.Tensor(
            [
                [
                    1,
                    2,
                    3,
                    3,
                    4,
                    10,
                    5,
                    4,
                    6,
                    7,
                    8,
                    7,
                    9,
                    10,
                    5,
                    4,
                    6,
                    7,
                    8,
                    7,
                    9,
                    10,
                    1,
                    2,
                    3,
                    3,
                    4,
                ],
                [
                    1,
                    2,
                    3,
                    3,
                    4,
                    10,
                    5,
                    4,
                    6,
                    7,
                    8,
                    7,
                    9,
                    10,
                    1,
                    2,
                    3,
                    3,
                    4,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                    0,
                ],
            ]
        ),
        torch.Tensor([1.0, 0.7037037037037037]),
        dict_int2lab,
        task="encode",
    )
    lens = (encoded_seq_pieces * encoded_seq_ids.shape[1]).round().int()
    # decode from torch tensors (batch, batch_lens)
    words_seq = spm(encoded_seq_ids, encoded_seq_pieces, task="decode")
    assert words_seq == gt, "output not the same"
    # decode from a list of bpe sequence (without padding)
    hyps_list = [
        encoded_seq_ids[0].int().tolist(),
        encoded_seq_ids[1][: lens[1]].int().tolist(),
    ]
    words_seq = spm(hyps_list, task="decode_from_list")
    assert words_seq == gt, "output not the same"

tests/unittests/test_tokenizer.py
TzuchengChang/NASS
"""Library for the HuggingFace (HF) repositories.

Authors
 * Mirco Ravanelli 2022
 * Andreas Nautsch 2022, 2023
"""
import os
import csv
from speechbrain.utils.data_utils import download_file
from tests.consistency.test_recipe import __skip_list


def run_HF_check(
    recipe_folder="tests/recipes", field="HF_repo", output_folder="tests/tmp",
):
    """Checks if the code reported in the readme files of the HF repository is
    runnable. Note: the tests run the code marked as python in the readme file.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder containing csv recipe files summarizing all the recipes in the repo.
    field: string
        Field of the csv recipe file containing the links to HF repos.
    output_folder: path
        Where to download the HF readme files.

    Returns
    ---------
    check: True
        True if all the code runs, False otherwise.
    """
    # Detect list of HF repositories
    HF_repos = repo_list(recipe_folder, field)

    # Set up output folder
    os.makedirs(output_folder, exist_ok=True)
    os.chdir(output_folder)

    # Checking all detected repos
    check = True
    for i, repo in enumerate(HF_repos):
        print("(%i/%i) Checking %s..." % (i + 1, len(HF_repos), repo))
        if not check_repo(repo):
            check = False
    return check


def repo_list(recipe_folder="tests/recipes", field="HF_repo"):
    """Get the list of HF recipes in the csv recipe file.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder containing csv recipe files summarizing all the recipes in the repo.
    field: string
        Field of the csv recipe file containing the links to HF repos.

    Returns
    ---------
    HF_repos: list
        List of the detected HF repos.
    """
    HF_repos = []

    # Loop over all recipe CSVs
    for recipe_csvfile in os.listdir(recipe_folder):
        if recipe_csvfile in __skip_list:
            continue

        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvf:
            reader = csv.DictReader(csvf, delimiter=",", skipinitialspace=True)
            for row in reader:
                if len(row[field]) > 0:
                    repos = row[field].split(" ")
                    for repo in repos:
                        HF_repos.append(repo)
    HF_repos = set(HF_repos)
    return HF_repos


def check_repo(HF_repo):
    """Runs the code reported in the README file of the given HF_repo. It checks
    if the code runs without errors.

    Arguments
    ---------
    HF_repo: string
        URL of the HF repository to check.

    Returns
    ---------
    check: bool
        True if all the code runs, False otherwise.
    """
    exp_name = os.path.basename(HF_repo)
    if HF_repo[-1] == "/":
        readme_file = HF_repo + "raw/main/README.md"
    else:
        readme_file = HF_repo + "/raw/main/README.md"

    dest_file = exp_name + ".md"
    download_file(readme_file, dest_file)

    code_snippets = []
    code = []
    flag = False
    check = True
    with open(dest_file, "r") as f:
        for line in f:
            if "```python" in line:
                flag = True
                code = []
            elif "```" in line and flag:
                flag = False
                code_snippets.append(code)
            elif flag:
                if len(line.strip()) > 0:
                    # adjust local audio paths 'tests/samples' -> '../samples'
                    if "tests/samples" in line:
                        line = line.replace("tests/samples", "../samples")
                    code.append(line)

    for code in code_snippets:
        try:
            exec("".join(code))
        except Exception as e:
            print("\t" + str(e))
            check = False
            print("\tERROR: cannot run code snippet in %s" % (HF_repo))
            print("---\n" + "".join(code) + "---\n")
    return check

tests/utils/check_HF_repo.py
TzuchengChang/NASS
"""This library contains functions that checks the dosctrings

Authors
 * Mirco Ravanelli 2022
"""

import re
from speechbrain.utils.data_utils import get_all_files


def extractName(s, search_class=False):
    """Extracts the names of the function or classes in the input string.

    Arguments
    ---------
    s: string
        Input string where to search for function or class names.
    search_clas: bool
        If True, searches for class names.

    Returns
    -------
    string: string
        Name of the function or class detected.
    """
    string = ""
    if search_class:
        regexp = re.compile(r"(class)\s(.*)\:")
    else:
        regexp = re.compile(r"(def)\s(.*)\(.*\)\:")
    for m in regexp.finditer(s):
        string += m.group(2)
    return string


def check_docstrings(
    base_folder=".", check_folders=["speechbrain", "tools", "templates"]
):
    """Checks if all the functions or classes have a docstring.

    Arguments
    ---------
    base_folder: path
        The main folder of speechbrain.
    check_folders: list
        List of subfolders to check.

    Returns
    -------
    check: bool
        True if all the functions/classes have a docstring, False otherwise.
    """
    # Search all python libraries in the folder of interest
    lib_lst = get_all_files(
        base_folder,
        match_and=[".py"],
        match_or=check_folders,
        exclude_or=[".pyc"],
    )
    check = True
    # Loop over the detected libraries
    for libpath in lib_lst:

        if "__" in libpath:
            continue
        print("Checking %s..." % (libpath))

        # Support variable initalization
        fun_name = libpath
        class_name = libpath
        check_line = True
        is_class = False
        first_line = True
        with open(libpath) as f:
            for line in f:

                # Remove spaces or tabs
                line = line.strip()

                # Avoid processing lines with the following patterns
                if ">>>" in line:
                    continue
                if "..." in line:
                    continue
                if len(line) == 0:
                    continue
                if line[0] == "#":
                    continue

                # Check if the docstring is written after the class/funct declaration
                if check_line:
                    if line[0] != '"' and not (is_class):
                        if line[0:2] == 'r"':
                            check_line = False
                            continue
                        check = False
                        if first_line:
                            print(
                                "\tERROR: The library %s must start with a docstring. "
                                % (libpath)
                                + "Please write it. For more info, see tests/consistency/DOCSTRINGS.md"
                            )
                        else:
                            print(
                                "\tERROR: The function %s in %s has no docstring. "
                                % (fun_name, libpath)
                                + "Please write it. For more info, see tests/consistency/DOCSTRINGS.md"
                            )
                    if line[0] != '"' and is_class:
                        if line[0:2] == 'r"':
                            check_line = False
                            continue

                        check = False
                        print(
                            "\tERROR: The class %s in %s has no docstring. "
                            % (class_name, libpath)
                            + "Please write it. For more info, see tests/consistency/DOCSTRINGS.md"
                        )

                    # Update support variables
                    check_line = False
                    is_class = False
                    first_line = False
                    continue

                # Extract function name (if any)
                fun_name = extractName(line)
                if len(fun_name) > 0 and line[0:3] == "def":
                    if fun_name[0] == "_":
                        continue
                    check_line = True

                # Extract class name (if any)
                class_name = extractName(line, search_class=True)
                if len(class_name) > 0 and line[0:5] == "class":
                    check_line = True
                    is_class = True
    return check

tests/utils/check_docstrings.py
TzuchengChang/NASS
"""Libraries for automatic finding URLs in the files and checking if they are
reachable.

Authors
 * Mirco Ravanelli 2022
"""
import os
import re
import time
import requests
from tqdm.contrib import tqdm
from speechbrain.utils.data_utils import get_all_files


def get_url(path):
    """This function searches for the URLs in the specified file.

    Arguments
    ---------
    path: path
        Path of the file where to search for URLs.

    Returns
    -------
    urls: list
        a list of all the URLs found in the specified path.
    """
    # Check if files exist
    if not (os.path.exists(path)):
        print("File %s not found!" % (path))
        return False

    # Read the file
    with open(path, "r") as file:
        text = file.read()

    # Set up Regex for URL detection
    url_regex = re.compile(
        r"((https?):((//)|(\\\\))+([\w\d:#@%/;$()~_?\+-=\\\.&](#!)?)*)",
        re.DOTALL,
    )
    urls = re.findall(url_regex, text)

    return list(set(urls))


def get_all_urls(file_lst, avoid_urls):
    """This function searches for all the URLs in the specified file list

    Arguments
    ---------
    file_lst: list
        List of the files where to search for URLs.
    avoid_urls: list
        List of URLs to avoid.

    Returns
    -------
    urls: dict
        A dictionary where the keys are the detected URLs and the values
    are the files where the URLs are found.
    """
    all_urls = {}

    for path in file_lst:
        if ".gz" in path:
            continue

        urls = get_url(path)

        for url in urls:

            # Clean up urls
            url = url[0].split(")")[0]
            if (
                url[-1] == "."
                or url[-1] == ","
                or url[-1] == " "
                or url[-1] == "/"
            ):
                url = url[:-1]

            if url in avoid_urls:
                continue

            if url not in all_urls:
                all_urls[url] = []
            all_urls[url].append(path)
    return all_urls


def check_url(url):
    """Cheks if an URL is broken

    Arguments
    ---------
    url: string
        URL to check

    Returns
    -------
    Bool
        False if the URL is broken, True otherwise.
    """
    try:
        response = requests.head(url)
        if response.status_code == 404 or response.status_code > 499:
            return False
        else:
            return True
    except requests.ConnectionError:
        return False


def check_links(
    folder=".",
    match_or=[".py", ".md", ".txt"],
    exclude_or=[".pyc"],
    avoid_files=[""],
    avoid_urls=["http:/", "http://", "https:/", "https://"],
):
    """This test checks if the files in the specified folders contain broken URLs

    Arguments
    ---------
    folder: path
        The top Folder for searching for the files.
    match_or: list
        Used to specify the extensions of the files to check.
    exclude_or: list
        Used to avoid some file extensions.
    avoid_files: list
        Used to avoid testing some specific file.
    """

    check_test = True
    # Find all the files that potentially contain urls
    file_lst = get_all_files(folder, match_or=match_or, exclude_or=exclude_or)

    # Get urls for the list of files - unique list
    all_urls = get_all_urls(file_lst, avoid_urls)

    # Check all the urls
    with tqdm(all_urls) as all_urls_progressbar:
        for url in all_urls_progressbar:
            time.sleep(1)
            if not (check_url(url)):
                check_test = False
                print("WARNING: %s is DOWN!" % (url))
                for path in all_urls[url]:
                    print("\t link detected in %s" % (path))
    return check_test

tests/utils/check_url.py
TzuchengChang/NASS
"""Tests for checking consistency between yaml files and their corresponding training scripts.

Authors
 * Mirco Ravanelli 2022
 * Andreas Nautsch 2022
"""

import os
import re


def get_yaml_var(hparam_file):
    """Extracts from the input yaml file (hparams_file) the list of variables that
    should be used in the script file.

    Arguments
    ---------
    hparam_file : path
        Path of the yaml file containing the hyperparameters.

    Returns
    -------
    var_list: list
        list of the variables declared in the yaml file (sub-variables are not
        included).
    """
    var_lst = []
    with open(hparam_file) as f:
        for line in f:
            # Avoid empty lists or comments - skip pretrainer definitions
            if (
                len(line) <= 1
                or line.startswith("#")
                or "speechbrain.utils.parameter_transfer.Pretrainer" in line
            ):
                continue

            # Remove trailing characters
            line = line.rstrip()

            # Check for variables (e.g., 'key:' or '- !ref')
            if line.find(":") != -1 or line.find("- !ref") != -1:
                var_name = line[: line.find(":")]
                # The variables to check are like "key:" (we do not need to check
                # subvariavles as " key:")
                if not (
                    var_name[0] == " "
                    or var_name[0] == "\t"
                    or "!!" in line
                    or "!apply" in line
                ):
                    var_lst.append(var_name)
                # Check for the reference pattern
                # note: output_folder: !ref results/<experiment_name>/<seed> pattern
                if line.find("!ref") != -1:
                    # Check for 'annotation_list_to_check: [!ref <train_csv>, !ref <valid_csv>]' pattern
                    for subline in line.split("<"):
                        sub_var = subline[: subline.find(">")]
                        # Check for 'models[generator]' pattern (dictionary reference)
                        dict_pos = sub_var.find("[")
                        if dict_pos != -1:
                            sub_var = sub_var[:dict_pos]
                        # Remove variables already used in yaml
                        if sub_var in var_lst:
                            var_lst.remove(sub_var)
    return var_lst


def detect_script_vars(script_file, var_lst):
    """Detects from the input script file (script_file) which of given variables (var_lst) are demanded.

    Arguments
    ---------
    script_file : path
        Path of the script file needing the hyperparameters.
    var_lst : list
        list of the variables declared in the yaml file.

    Returns
    -------
    detected_var: list
        list of the variables detected in the script file.
    """
    var_types = [
        "hparams.",
        "modules.",
        'attr(self.hparams, "',
        'hparams.get("',
    ]
    detected_var = []
    with open(script_file) as f:
        for line in f:
            for var in var_lst:
                # The pattern can be ["key"] or ".key"
                if '["' + var + '"]' in line:
                    if var not in detected_var:
                        detected_var.append(var)
                        continue  # no need to go through the other cases for this var
                # case: hparams[f"{dataset}_annotation"] - only that structure at the moment
                re_match = re.search(r"\[f.\{.*\}(.*).\]", line)
                if re_match is not None:
                    if re_match.group(1) in var:
                        print(
                            "\t\tWARNING: potential inconsistency %s maybe used in %s (or not)."
                            % (var, re_match.group(0))
                        )
                        if var not in detected_var:
                            detected_var.append(var)
                            continue

                # case: hparams[f"annotation_{dataset}"] - only that structure at the moment
                re_match = re.search(r"\[f.(.*)\{.*\}.\]", line)
                if re_match is not None:
                    if re_match.group(1) in var:
                        print(
                            "\t\tWARNING: potential inconsistency %s maybe used in %s (or not)."
                            % (var, re_match.group(0))
                        )
                        if var not in detected_var:
                            detected_var.append(var)
                            continue

                # Chek var types
                # Chek var types
                for var_type in var_types:
                    if var_type + var in line:
                        if var not in detected_var:
                            detected_var.append(var)
                            continue
                # case: tea_enc_list.append(hparams['tea{}_enc'])
                re_var = re.search(r"\[.(.*){}(.*).\]", line)
                if re_var is not None:
                    re_var_pattern = re_var.group(1) + ".*" + re_var.group(2)
                    re_pattern = re.search(re_var_pattern, var)
                    if re_pattern is not None:
                        if re_pattern.group() == var:
                            detected_var.append(var)
                            continue
    return detected_var


def check_yaml_vs_script(hparam_file, script_file):
    """Checks consistency between the given yaml file (hparams_file) and the
    script file. The function detects if there are variables declared in the yaml
    file, but not used in the script file.

    Arguments
    ---------
    hparam_file : path
        Path of the yaml file containing the hyperparameters.
    script_file : path
        Path of the script file (.py) containing the training recipe.

    Returns
    -------
    Bool
        This function returns False if some mismatch is detected and True otherwise.
        An error is raised to inform about which variable has been declared but
        not used.
    """
    print("Checking %s..." % (hparam_file))

    # Check if files exist
    if not (os.path.exists(hparam_file)):
        print("File %s not found!" % (hparam_file,))
        return False

    if not (os.path.exists(script_file)):
        print("File %s not found!" % (script_file,))
        return False

    # Getting list of variables declared in yaml
    var_lst = get_yaml_var(hparam_file)

    # Detect which of these variables are used in the script file
    detected_vars_train = detect_script_vars(script_file, var_lst)

    # Check which variables are declared but not used
    default_run_opt_keys = [
        "debug",
        "debug_batches",
        "debug_epochs",
        "device",
        "cpu",
        "data_parallel_backend",
        "distributed_launch",
        "distributed_backend",
        "find_unused_parameters",
        "jit_module_keys",
        "auto_mix_prec",
        "max_grad_norm",
        "nonfinite_patience",
        "noprogressbar",
        "ckpt_interval_minutes",
        "grad_accumulation_factor",
        "optimizer_step_limit",
    ]
    unused_vars = list(
        set(var_lst) - set(detected_vars_train) - set(default_run_opt_keys)
    )
    for unused_var in unused_vars:
        print(
            '\tERROR: variable "%s" not used in %s!' % (unused_var, script_file)
        )

    return len(unused_vars) == 0


def extract_patterns(lines, start_pattern, end_pattern):
    """Extracts a variables from start_pattern to end_pattern.

    Arguments
    ---------
    lines: list
        List of strings to parse.
    start_pattern: string
        String that indicated the start of the pattern.
    end_pattern: string
        String that indicated the end of the pattern.

    Returns
    -------
    var_lst: list
        List of variables detected.
    """
    var_lst = []

    for line in lines:
        start_indexes = [
            index
            for index in range(len(line))
            if line.startswith(start_pattern, index)
        ]
        for index in start_indexes:
            start_var = index + len(start_pattern)
            line_src = line[start_var:]
            var_name = ""
            for char in line_src:
                if char in end_pattern:
                    break
                var_name = var_name + char
            var_lst.append(var_name)
    return var_lst


def check_module_vars(
    hparam_file, script_file, module_key="modules:", module_var="self.modules."
):
    """Checks if the variables self.moduled.var are properly declared in the
    hparam file.

    Arguments
    ---------
    hparam_file : path
        Path of the yaml file containing the hyperparameters.
    script_file : path
        Path of the script file (.py) containing the training recipe.
    module_key: string
        String that denoted the start of the module in the hparam file.
    module_var: string
        String that denoted the start of the module in the script file.
    Returns
    -------
    Bool
        This function returns False if some mismatch is detected and True otherwise.
        An error is raised to inform about which variable has been used but
        not declared.
    """
    stop_char = [" ", ",", "(", ")", "[", "]", "{", "}", ".", ":", "\n"]
    module_block = False
    end_block = [" ", "\t"]
    avoid_lst = ["parameters", "keys", "eval", "train"]

    # Extract Modules variables from the hparam file
    module_vars_hparams = []
    with open(hparam_file) as f:
        for line in f:
            if module_key in line:
                module_block = True
                continue
            if line[0] not in end_block:
                module_block = False

            if module_block:
                var = line.strip().split(":")[0]
                module_vars_hparams.append(var)

    # Extract Modules variables from the script file
    with open(script_file) as file:
        lines = file.readlines()
        lines = [line.rstrip() for line in lines]

    module_var_script = extract_patterns(lines, module_var, stop_char)
    module_var_script = set(module_var_script)
    for avoid in avoid_lst:
        if avoid in module_var_script:
            module_var_script.remove(avoid)

    # Remove optional variables "if hasattr(self.modules, "env_corrupt"):"
    stop_char.append('"')
    opt_vars = extract_patterns(lines, 'if hasattr(self.modules, "', stop_char)
    opt_vars.extend(
        extract_patterns(lines, 'if hasattr(self.hparams, "', stop_char)
    )

    # Remove optional
    for avoid in set(opt_vars):
        if avoid in module_var_script:
            module_var_script.remove(avoid)

    # Check Module variavles
    unused_vars = list(set(module_var_script) - set(module_vars_hparams))

    for unused_var in unused_vars:
        print(
            '\tERROR: variable "self.modules.%s" used in %s, but not listed in the "modules:" section of %s'
            % (unused_var, script_file, hparam_file)
        )
    return len(unused_vars) == 0

tests/utils/check_yaml.py
TzuchengChang/NASS
"""Library for running recipe tests.

Authors
 * Mirco Ravanelli 2022
 * Andreas Nautsch 2022, 2023
"""
import os
import re
import csv
import sys
import pydoc
from time import time
import subprocess as sp
from hyperpyyaml import load_hyperpyyaml
from tests.consistency.test_recipe import __skip_list


def check_row_for_test(row, filters_fields, filters, test_field):
    """Checks if the current row of the csv recipe file has a test to run.

    Arguments
    ---------
    row: dict
        Line of the csv file (in dict from).
    filters_fields: list
        This can be used with the "filter" variable
        to run only some tests. For instance, filters_fields=['Task'] and filters=['ASR'])
        will only run tests for ASR recipes.
    filters: list
        See above.
    test_field: string
        Key of the input dictionary that contains the test flags.

    Returns
    ---------
    test: bool
        True if the line must be tested, False otherwise.
    """
    test = True
    for i, field in enumerate(filters_fields):
        field_values = filters[i]
        if type(field_values) == str:
            # ... AND ... filter
            if not (field_values == row[field]):
                test = False
        elif type(field_values) == list:  # type(field) == list
            # ... AND (... OR ...) ... filter; at least one entry of the list matches
            test_flag = False
            for filt in field_values:
                if filt == row[field]:
                    test_flag = True
            test = test and test_flag
        else:
            print("\tError in filters_fields and filters definition.")
            test = False

    if test:
        test_flag = row[test_field].strip()
        if len(test_flag) == 0:
            test = False
    return test


def prepare_test(
    recipe_folder="tests/recipes",
    script_field="Script_file",
    hparam_field="Hparam_file",
    test_field="test_debug_flags",
    check_field="test_debug_checks",
    filters_fields=[],
    filters=[],
):
    """Extracts all the needed information to run the recipe test.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder containing csv recipe files summarizing all the recipes in the repo.
    script_field: str
        Field of the csv recipe file containing the path of the script to run.
    hparam_field: str
        Field of the csv recipe file containing the path of the hparam file.
    test_field: string
        Field of the csv recipe file containing the test flags.
    check_field: string
        Field of the csv recipe file containing the checks to perform.
    filters_fields: list
        This can be used with the "filter" variable
        to run only some tests. For instance, filters_fileds=['Task'] and filters=['ASR'])
        will only run tests for ASR recipes.
    filters: list
        See above.

    Returns
    ---------
    test_script: dict
        A Dictionary containing recipe IDs as keys and test_scripts as values.
    test_hparam: dict
        A dictionary containing recipe IDs as keys and hparams as values.
    test_flag: dict
        A dictionary containing recipe IDs as keys and the test flags as values.
    test_check: dict
        A dictionary containing recipe IDs as keys and the checks as values.
    """

    # Dictionary initialization
    test_script = {}
    test_hparam = {}
    test_flag = {}
    test_check = {}

    # Loop over all recipe CSVs
    print(f"\tfilters_fields={filters_fields} => filters={filters}")
    for recipe_csvfile in os.listdir(recipe_folder):
        # skip setup scripts; consider CSV files only
        if recipe_csvfile in __skip_list:
            continue

        print(f"Loading recipes from: {recipe_csvfile}")
        # Detect needed information for the recipe tests
        with open(
            os.path.join(recipe_folder, recipe_csvfile), newline=""
        ) as csvf:
            reader = csv.DictReader(csvf, delimiter=",", skipinitialspace=True)
            for row_id, row in enumerate(reader):
                recipe_id = f"{recipe_csvfile[:-4]}_row_{row_id+2}"
                if not (
                    check_row_for_test(row, filters_fields, filters, test_field)
                ):
                    print(
                        f"\tSkipped {recipe_id} - lacking test_field={test_field}"
                    )
                    continue
                test_script[recipe_id] = row[script_field].strip()
                test_hparam[recipe_id] = row[hparam_field].strip()
                test_flag[recipe_id] = row[test_field].strip()
                test_check[recipe_id] = row[check_field].strip()

    return test_script, test_hparam, test_flag, test_check


def check_files(
    check_str, output_folder, recipe_id, pattern=r"file_exists=\[(.*?)\]"
):
    """Checks if the output folder created by the test has the expected files.

    Arguments
    ---------
    check_str: str
        String summarizing the checks to perform.
    output_folder: path
        The path where to check the files.
    recipe_id: str
        Unique ID of the recipe.
    pattern: str
        The pattern used to extract the list of files to check from check_str.

    Returns
    ---------
    check: bool
        True if all the files are found, False otherwise.
    """

    check = True
    files_to_check = re.search(pattern, check_str)
    files_to_check = files_to_check.group(1).split(",")

    for file_to_check in files_to_check:
        check_path = os.path.join(output_folder, file_to_check)
        if not (os.path.exists(check_path)):
            print(
                "\tERROR: The recipe %s does not contain the expected file %s"
                % (recipe_id, check_path)
            )
            check = False
    return check


def check_performance(
    check_str, output_folder, recipe_id, pattern=r"performance_check=\[(.*?)\]"
):
    """Checks if the performance achieved by the recipe matches with the
    expectations. This is done by adding a performance_check entry in the recipe
    check field of the csv recipe file
    For instance: performance_check=[train_log.txt, train loss, <=15, epoch: 2]),
    will check the variable "train_loss" in the train_log.txt at epoch 2. It will
    raise an error if the train_loss is >15.

    Arguments
    ---------
    check_str: str
        String summarizing the checks to perform.
    output_folder: path
        The path where the recipe files are stored.
    recipe_id: str
        Unique ID of the recipe.
    pattern: str
        The pattern used to extract the list of files to check from check_str.

    Returns
    ---------
    check: bool
        True if all the files are found, False otherwise.
    """

    check = True
    performance_to_check = re.search(pattern, check_str)
    if performance_to_check is None:
        return check

    # Getting the needed information from the "performance_check" entry
    performance_to_check = performance_to_check.group(1).split(",")
    filename = performance_to_check[0].strip()
    filename = os.path.join(output_folder, filename)
    variable = performance_to_check[1].strip()
    threshold = performance_to_check[2].strip()
    epoch = performance_to_check[3].strip()

    if not (os.path.exists(filename)):
        print(
            "\tERROR: The file %s of recipe %s does not exist (needed for performance checks)"
            % (filename, recipe_id)
        )

        return False

    # Real all the lines of the performance file
    with open(filename) as file:
        lines = file.readlines()

    # Fitler the lines
    lines_filt = []
    for line in lines:
        if epoch in line:
            lines_filt.append(line)

    # Raising an error if there are no lines after applying the filter
    if len(lines_filt) == 0:
        print(
            "\tERROR: No entries %s in %s (recipe %s). See performance_check entry."
            % (epoch, filename, recipe_id)
        )
        return False

    for line in lines_filt:

        # Search variable value
        pattern = variable + ": " + "(.*?) "
        var_value = re.search(pattern, line)

        if var_value is None:
            print(
                "\tERROR: The file %s of recipe %s does not contain the variable %s (needed for performance checks)"
                % (filename, recipe_id, variable)
            )
            return False
        var_value = float(var_value.group(1))
        check = check_threshold(threshold, var_value)

        if not (check):
            print(
                "\tERROR: The variable %s of file %s (recipe %s) violated the specified threshold (%s %s)"
                % (variable, filename, recipe_id, var_value, threshold)
            )

        break

    return check


def check_threshold(threshold, value):
    """Checks if the value satisfied the threshold constraints.

    Arguments
    ---------
    threshold: str
        String that contains the contains. E.g, ">=10" or "==15" or "<5".
    value: float
        Float corresponding to the value to test

    Returns
    ---------
    bool
        True if the constraint is satisfied, False otherwise.
    """

    # Get threshold value:
    th_value = float(
        threshold.strip().replace("=", "").replace(">", "").replace("<", "")
    )

    # Check Threshold
    if "==" in threshold:
        return value == th_value

    elif ">=" in threshold:
        return value >= th_value

    elif ">" in threshold:
        return value > th_value

    elif "<=" in threshold:
        return value <= th_value

    elif "<" in threshold:
        return value < th_value
    else:
        return False


def run_test_cmd(cmd, stdout_file, stderr_file):
    """Runs the command corresponding to a recipe test. The standard output and
    the standard error is saved in the specified paths.

    Arguments
    ---------
    cmd: str
        String corresponding to the command to run.
    stdout_file: path
        File where standard output is stored.
    stderr_file: path
        File where standard error is stored.

    Returns
    ---------
    rc: bool
        The return code obtained after running the command. If 0, the test is
        run without errors. If >0 the execution failed.
    """
    f_stdout = open(stdout_file, "w")
    f_stderr = open(stderr_file, "w")
    child = sp.Popen([cmd], stdout=f_stdout, stderr=f_stderr, shell=True)
    child.communicate()[0]
    rc = child.returncode
    f_stdout.close()
    f_stderr.close()
    return rc


def run_recipe_tests(
    recipe_folder="tests/recipes",
    script_field="Script_file",
    hparam_field="Hparam_file",
    test_field="test_debug_flags",
    check_field="test_debug_checks",
    run_opts="--device=cpu",
    output_folder="tests/tmp/",
    filters_fields=[],
    filters=[],
    do_checks=True,
    download_only=False,
    run_tests_with_checks_only=False,
):
    """Runs the recipes tests.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder containing csv recipe files summarizing all the recipes in the repo.
    script_field: str
        Field of the csv recipe file containing the path of the script to run.
    hparam_field: str
        Field of the csv recipe file containing the path of the hparam file.
    test_field: string
        Field of the csv recipe file containing the test flags.
    check_field: string
        Field of the csv recipe file containing the checks to perform.
    run_opts: string
        Additional flags to add for the tests (see run_opts of speechbrain/core.py).
    output_folder: string
        Folder where the output of the tests are saved.
    filters_fields: list
        This can be used with the "filter" variable
        to run only some tests. For instance, filters_fileds=['Task'] and filters=['ASR'])
        will only run tests for ASR recipes.
    filters: list
        See above.
    do_checks: bool (default: True)
        If True performs the checks on the output folder (when the check_field is not empty).
    download_only: bool (default: False)
        If True skips running/checking tests after downloading relevant pre-trained data (prepare for offline testing).
    run_tests_with_checks_only: bool (default: False)
        If True skips all tests that do not have performance check criteria defined.

    Returns
    -------
    check: True
        True if all the recipe tests pass, False otherwise.

    Example
    -------
    python -c 'from speechbrain.utils.recipe_tests import run_recipe_tests; print("TEST FAILED!") if not(run_recipe_tests(filters_fields=["Dataset", "Task"], filters=[["AISHELL-1", "CommonVoice"], "SSL"])) else print("TEST PASSED")'
    """
    # Create the output folder (where the tests results will be saved)
    os.makedirs(output_folder, exist_ok=True)
    print("Test ouputs will be put in %s" % (output_folder))

    # Read the csv recipe file and detect which tests we have to run
    test_script, test_hparam, test_flag, test_check = prepare_test(
        recipe_folder,
        script_field,
        hparam_field,
        test_field=test_field,
        check_field=check_field,
        filters_fields=filters_fields,
        filters=filters,
    )

    # Early stop if there are no recipes to test
    if len(test_script) == 0:
        print("No recipes found for testing (please check recipe filters).")
        return False

    # Download all upfront
    if download_only:
        for i, recipe_id in enumerate(test_script.keys()):
            # If we are interested in performance checks only, skip
            check_str = test_check[recipe_id].strip()
            if run_tests_with_checks_only:
                if len(check_str) == 0:
                    continue

            print(
                "(%i/%i) Collecting pretrained models for %s..."
                % (i + 1, len(test_script.keys()), recipe_id)
            )

            output_fold = os.path.join(output_folder, recipe_id)
            os.makedirs(output_fold, exist_ok=True)
            stdout_file = os.path.join(output_fold, "stdout.txt")
            stderr_file = os.path.join(output_fold, "stderr.txt")

            cmd = (
                "python -c 'import sys;from hyperpyyaml import load_hyperpyyaml;import speechbrain;"
                "hparams_file, run_opts, overrides = speechbrain.parse_arguments(sys.argv[1:]);"
                "fin=open(hparams_file);hparams = load_hyperpyyaml(fin, overrides);fin.close();"
                # 'speechbrain.create_experiment_directory(experiment_directory=hparams["output_folder"],'
                # 'hyperparams_to_save=hparams_file,overrides=overrides,);'
            )
            with open(test_hparam[recipe_id]) as hparam_file:
                for line in hparam_file:
                    if "pretrainer" in line:
                        cmd += 'hparams["pretrainer"].collect_files();hparams["pretrainer"].load_collected(device="cpu");'
                    elif "from_pretrained" in line:
                        field = line.split(":")[0].strip()
                        cmd += f'hparams["{field}"]'
            cmd += (
                "' "
                + test_hparam[recipe_id]
                + " --output_folder="
                + output_fold
                + " "
                + test_flag[recipe_id]
                + " "
                + run_opts
            )

            # Prepare the test
            run_test_cmd(cmd, stdout_file, stderr_file)

        return False

    # Run  script (check how to get std out, std err and save them in files)
    check = True
    for i, recipe_id in enumerate(test_script.keys()):
        # If we are interested in performance checks only, skip
        check_str = test_check[recipe_id].strip()
        if run_tests_with_checks_only:
            if len(check_str) == 0:
                continue

        print(
            "(%i/%i) Running test for %s..."
            % (i + 1, len(test_script.keys()), recipe_id)
        )

        output_fold = os.path.join(output_folder, recipe_id)
        os.makedirs(output_fold, exist_ok=True)
        stdout_file = os.path.join(output_fold, "stdout.txt")
        stderr_file = os.path.join(output_fold, "stderr.txt")

        # Check for setup scripts
        setup_script = os.path.join(
            "tests/recipes/setup",
            test_script[recipe_id][:-3].replace("/", "_"),
            test_hparam[recipe_id]
            .replace(os.path.dirname(test_script[recipe_id]), "")[1:-5]
            .replace("/", "_"),
        )
        if os.path.exists(setup_script):
            os.system(setup_script)

        # Composing command to run
        cmd = (
            f"PYTHONPATH={os.getcwd() + '/' + os.path.dirname(test_script[recipe_id])} python "
            + test_script[recipe_id]
            + " "
            + test_hparam[recipe_id]
            + " --output_folder="
            + output_fold
            + " "
            + test_flag[recipe_id]
            + " "
            + run_opts
        )

        # add --debug if no do_checks to save testing time
        if not do_checks:
            cmd += " --debug --debug_persistently"

        # Running the test
        time_start = time()
        return_code = run_test_cmd(cmd, stdout_file, stderr_file)
        test_duration = time() - time_start
        print("\t... %.2fs" % test_duration)

        # Tear down
        td_script = os.path.join(os.path.dirname(setup_script), "tear_down")
        if os.path.exists(td_script):
            os.system(td_script)

        # Check return code
        if return_code != 0:
            print(
                "\tERROR: Error in %s (%s). Check %s and %s for more info."
                % (recipe_id, test_hparam[recipe_id], stderr_file, stdout_file)
            )
            check = False

        # Checks
        if do_checks and len(check_str) > 0:
            print("\t...checking files & performance...")

            # Check if the expected files exist
            check &= check_files(check_str, output_fold, recipe_id)
            check &= check_performance(check_str, output_fold, recipe_id)

    return check


def load_yaml_test(
    recipe_folder="tests/recipes",
    script_field="Script_file",
    hparam_field="Hparam_file",
    test_field="Hparam_file",
    filters_fields=[],
    filters=[],
    avoid_list=[],
    rir_folder="tests/tmp/rir",
    data_folder="tests/tmp/yaml",
    output_folder="tests/tmp/",
):
    """Tests if the yaml files can be loaded without errors.

    Arguments
    ---------
    recipe_folder: path
        Path of the folder containing csv recipe files summarizing all the recipes in the repo.
    script_field: str
        Field of the csv recipe file containing the path of the script to run.
    hparam_field: str
        Field of the csv recipe file containing the path of the hparam file.
    test_field: string
        Field of the csv recipe file containing the test flags.
    filters_fields: list
        This can be used with the "filter" variable
        to run only some tests. For instance, filters_fileds=['Task'] and filters=['ASR'])
        will only run tests for ASR recipes.
    filters: list
        See above.
    avoid_list: list
        List of hparam file not to check.
    rir_folder:
        This overrides the rir_folder; rir_path, and openrir_folder usually specified in the hparam files.
    data_folder:
        This overrides the data_folder usually specified in the hparam files.
    output_folder:
        This overrides the output_folder usually specified in the hparam files.

    Returns
    ---------
    check: True
        True if all the hparam files are loaded correctly, False otherwise.
    """

    # Get current working directory
    cwd = os.getcwd()

    # Set data_foler and output folder
    data_folder = os.path.join(cwd, data_folder)
    output_folder = os.path.join(cwd, output_folder)
    rir_folder = os.path.join(cwd, rir_folder)

    # Additional overrides
    add_overrides = {
        "manual_annot_folder": data_folder,
        "musan_folder": data_folder,
        "tea_models_dir": data_folder,
        "wsj_root": data_folder,
        "tokenizer_file": data_folder,
        "commonlanguage_folder": data_folder,
        "tea_infer_dir": data_folder,
        "original_data_folder": data_folder,
        "pretrain_st_dir": data_folder,
        # RIR folder specifications -> all point to the same zip file: one download destination
        "rir_path": rir_folder,
        "rir_folder": rir_folder,
        "openrir_folder": rir_folder,
        "open_rir_folder": rir_folder,
        "data_folder_rirs": rir_folder,
    }

    # Additional overrides when extra !PLACEHOLDER are encountered (not: data_folder - output_folder)
    add_placeholder_overrides = {
        "wav2vec2_hub": "facebook/wav2vec2-large-960h-lv60-self",  # this might not hold for all set-ups
        "root_data_folder": data_folder,
        "wav2vec2_folder": f"{output_folder}/wav2vec2_checkpoint",
        # these will need refactoring at some point (recipe-depending values)
        "pretrained_tokenizer_path": "speechbrain/asr-wav2vec2-switchboard",
        "pretrained_lm_tokenizer_path": "speechbrain/asr-transformer-switchboard",
        "channels_path": None,
        "concepts_path": None,
    }

    # Read the csv recipe file and detect which tests we have to run
    test_script, test_hparam, test_flag, test_check = prepare_test(
        recipe_folder,
        script_field,
        hparam_field,
        test_field=test_field,
        filters_fields=filters_fields,
        filters=filters,
    )

    check = True
    for i, recipe_id in enumerate(test_script.keys()):
        hparam_file = test_hparam[recipe_id]
        script_file = test_script[recipe_id]

        # Changing working folder to recipe folder (as 'run_folder' to avoid name conflict with arg 'recipe_folder')
        script_folder = os.path.dirname(script_file)
        run_folder = os.path.join(cwd, script_folder)
        os.chdir(run_folder)

        # Avoid files lister in avoid_list
        if hparam_file in avoid_list:
            print(
                f"\t({i + 1}/{len(test_script.keys())}) Skipped: {hparam_file}! (check avoid_list for details)"
            )
            continue

        print(
            "(%i/%i) Checking %s..."
            % (i + 1, len(test_script.keys()), hparam_file)
        )

        # Get absolute path to the hparam file
        hparam_file = os.path.join(cwd, hparam_file)

        # Load hyperparameters file with command-line overrides
        overrides = {"data_folder": data_folder, "output_folder": output_folder}

        tag_custom_model = None
        # Append additional overrides when needed
        with open(hparam_file) as f:
            for line in f:
                # os.chdir(run_folder) is not changing sys.module, and pydoc.locate (in load_hyperpyyaml) fails
                if "new:custom_model" in line:
                    tag_custom_model = "custom_model"
                    custom_model_from_root = f"{script_folder.replace(os.sep, '.')}.{tag_custom_model}"
                    if pydoc.locate(custom_model_from_root) is not None:
                        sys.modules[tag_custom_model] = sys.modules[
                            custom_model_from_root
                        ]
                # check for !PLACEHOLDER overrides
                flag_continue = False
                for key, value in add_placeholder_overrides.items():
                    placeholder_pattern = key + ": !PLACEHOLDER"
                    if (
                        placeholder_pattern in line
                        and line.find(placeholder_pattern) == 0
                    ):
                        overrides.update({key: value})
                        flag_continue = True

                # if !PLACEHOLDER was substituted already, skip further pattern overrides for this line
                if flag_continue:
                    continue

                # check for additional overrides
                for key, value in add_overrides.items():
                    pattern = key + ":"
                    if pattern in line and line.find(pattern) == 0:
                        overrides.update({key: value})

        with open(hparam_file) as fin:
            try:
                _ = load_hyperpyyaml(fin, overrides)
            except Exception as e:
                print("\t" + str(e))
                check = False
                print("\tERROR: cannot load %s" % (hparam_file))
        if tag_custom_model is not None:
            if tag_custom_model in sys.modules:
                del sys.modules[tag_custom_model]
            for tcm_key in list(sys.modules.keys()):
                if tcm_key.startswith("recipes"):
                    del sys.modules[tcm_key]
    return check

tests/utils/recipe_tests.py
TzuchengChang/NASS
#!/usr/bin/env/python3
"""This is a test script for creating a list of expected outcomes (before refactoring);
then, manual editing might change YAMLs and/or code; another test runs to compare results
(after refactoring to before). The target is a list of known HF repos.

The goal is to identify to which extent changes break existing functionality.
Then, larger changes to code base can be rolled out more assured.

Authors
 * Andreas Nautsch, 2022, 2023
"""

import os
import sys
from tqdm import tqdm
import yaml
import torch  # noqa
import importlib  # noqa
import subprocess
import speechbrain  # noqa
from glob import glob
from copy import deepcopy
from torch.utils.data import DataLoader
from hyperpyyaml import load_hyperpyyaml
from speechbrain.utils.distributed import run_on_main  # noqa
from speechbrain.utils.train_logger import FileTrainLogger
from speechbrain.pretrained.interfaces import foreign_class  # noqa
from speechbrain.dataio.dataloader import LoopedLoader, make_dataloader


def init(
    new_interfaces_git="https://github.com/speechbrain/speechbrain",
    new_interfaces_branch="hf-interface-testing",
    new_interfaces_local_dir="tests/tmp/hf_interfaces",
):
    """Initialises a PR branch to: https://github.com/speechbrain/speechbrain/tree/hf-interface-testing

    Skip if the path as of `new_interfaces_local_dir` exists (e.g. by DIY init instead of via this script).

    Parameters
    ----------
    new_interfaces_git: str
        Your git repo (or default: `https://github.com/speechbrain/speechbrain`);
        can be specified in tests/utils/overrides.yaml
    new_interfaces_branch: str
        Default is `hf-interface-testing` (a git branch); can be specified in tests/utils/overrides.yaml
    new_interfaces_local_dir: str
        Default is `tests/tmp/hf_interfaces` (a local path); can be specified in tests/utils/overrides.yaml

    Returns
    -------
    str
        Local path of `updates_pretrained_models` where the update HF yaml/interface files can be found.
    """
    # set up git etc
    if not os.path.exists(new_interfaces_local_dir):
        # note: not checking for anything, whether it exists or not - or if there is a previous one already
        # clone repo with PR on updates_pretrained_models into local folder
        cmd_out_clone = subprocess.run(
            ["git", "clone", new_interfaces_git, new_interfaces_local_dir],
            capture_output=True,
        )
        print(f"\tgit clone log: {cmd_out_clone}")

        # cd into that local folder, switch branch to the one containing updates_pretrained_models & cd back
        cwd = os.getcwd()
        os.chdir(new_interfaces_local_dir)
        cmd_out_co = subprocess.run(
            ["git", "checkout", new_interfaces_branch], capture_output=True
        )
        print(f"\tgit checkout log: {cmd_out_co}")
        os.chdir(cwd)

    # return the valid local path with updates_pretrained_models
    updates_dir = f"{new_interfaces_local_dir}/updates_pretrained_models"
    return updates_dir


def get_model(repo, values, updates_dir=None, run_opts=None):
    """Fetches a pretrained model with the option the re-specify its hyperparameters & interface.

    Parameters
    ----------
    repo: str
        Source of pretrained model (assuming its within the HF speechbrain collection).
    values: dict
        Interface specification.
        Example: speechbrain:hf-interface-testing/updates_pretrained_models/ssl-wav2vec2-base-librispeech/test.yaml
    updates_dir: str
        Local folder with yaml:interface updates; None (default) = take original yaml/interface specification.
    run_opts: dict
        Run options, such as device

    Returns
    -------
    A pretrained model with a speechbrain.pretrained.interface or a custom interface.
    """
    # get the pretrained class; model & predictions
    kwargs = {
        "source": f"speechbrain/{repo}",
        "savedir": f"pretrained_models/{repo}",
    }

    # adjust symlinks
    hparams = f"pretrained_models/{repo}/hyperparams.yaml"
    if (
        "foreign" in values.keys()
    ):  # it's a custom model which has its own Python filename
        custom = f'pretrained_models/{repo}/{values["foreign"]}'
    # prepare model loading: is it the old -or- the new yaml/interface?
    if updates_dir is not None:
        # testing the refactoring; assuming all model data has been loaded already
        kwargs["source"] = f"{updates_dir}/{repo}"
        os.unlink(hparams)
        os.symlink(f"{updates_dir}/{repo}/hyperparams.yaml", hparams)
        if "foreign" in values.keys():
            os.unlink(custom)
            os.symlink(
                f'{updates_dir}/{repo}/{values["foreign"]}', custom,
            )
    else:
        # re:testing on develop? => simply unlink anything before and re:link from cached HF hub
        if os.path.exists(hparams):
            os.unlink(hparams)
        if "foreign" in values.keys():
            if os.path.exists(custom):
                os.unlink(custom)

    if run_opts is not None:
        kwargs["run_opts"] = run_opts

    print(f"\trepo: {repo}")
    # load pretrained model either via specified pretrained class or custom interface
    if "foreign" not in values.keys():
        print(f'\tspeechbrain.pretrained.{values["cls"]}')
        print(f"\tobj.from_hparams({kwargs})")
        obj = eval(f'speechbrain.pretrained.{values["cls"]}')
        model = obj.from_hparams(**kwargs)
    else:
        kwargs["pymodule_file"] = values["foreign"]
        kwargs["classname"] = values["cls"]
        model = foreign_class(**kwargs)

    return model


def get_prediction(repo, values, updates_dir=None):
    """Gets the prediction for one predefined audio example, pattern: {repo}/{values["sample"]} (see HF model card).

    Parameters
    ----------
    repo: str
        Source of pretrained model (assuming its within the HF speechbrain collection).
    values: dict
        Interface specification.
        Examples: speechbrain:hf-interface-testing/updates_pretrained_models/ssl-wav2vec2-base-librispeech/test.yaml
                  speechbrain:hf-interface-testing/updates_pretrained_models/asr-wav2vec2-librispeech/test.yaml
    updates_dir: str
        Controls whether/not we are in the refactored results (None: expected results; before refactoring).

    Returns
    -------
    Cleaned-up prediction results for yaml output (result logging & comparison through yaml de/serialization).
    """

    def sanitize(data):
        # cleanup data for yaml output (w/o this, yaml will make attempts to save torch/numpy arrays in their format)
        if isinstance(data, torch.Tensor):
            data = data.detach().cpu().numpy()
            if data.ndim:
                data = list(data)
        return data

    # get the pretrained model (before/after yaml/interface update)
    model = get_model(repo=repo, values=values, updates_dir=updates_dir)  # noqa

    try:
        # simulate batch from single file
        prediction = eval(
            f'model.{values["fnx"]}(model.load_audio("{repo}/{values["sample"]}", savedir="pretrained_models/{repo}").unsqueeze(0), torch.tensor([1.0]))'
        )

    except Exception:
        # use an example audio if no audio can be loaded
        print(f'\tWARNING - no audio found on HF: {repo}/{values["sample"]}')
        prediction = eval(
            f'model.{values["fnx"]}(model.load_audio("tests/samples/single-mic/example1.wav", savedir="pretrained_models/{repo}").unsqueeze(0), torch.tensor([1.0]))'
        )

    finally:
        del model

    return [sanitize(x[0]) for x in prediction]


def gather_expected_results(
    glob_filter="*",
    new_interfaces_git="https://github.com/speechbrain/speechbrain",
    new_interfaces_branch="hf-interface-testing",
    new_interfaces_local_dir="tests/tmp/hf_interfaces",
    yaml_path="tests/tmp/refactoring_results.yaml",
):
    """Before refactoring HF YAMLs and/or code, gather prediction results.

    Parameters
    ----------
    glob_filter: str
        Filter for a repo subset or a specific repo.
    new_interfaces_git: str
        Your git repo (or default: `https://github.com/speechbrain/speechbrain`);
        can be specified in tests/utils/overrides.yaml
    new_interfaces_branch: str
        Default is `hf-interface-testing` (a git branch); can be specified in tests/utils/overrides.yaml
    new_interfaces_local_dir: str
        Default is `tests/tmp/hf_interfaces` (a local path); can be specified in tests/utils/overrides.yaml
    yaml_path : str
        Path where to store/load refactoring testing results for later comparison.

    """
    # load results, if existing -or- new from scratch
    if os.path.exists(yaml_path):
        with open(yaml_path) as yaml_in:
            results = yaml.safe_load(yaml_in)
    else:
        results = {}

    # go through each repo
    updates_dir = init(
        new_interfaces_git, new_interfaces_branch, new_interfaces_local_dir
    )
    repos = map(os.path.basename, glob(f"{updates_dir}/{glob_filter}"),)
    for repo in repos:
        # skip if results are there
        if repo not in results.keys():
            # get values
            with open(f"{updates_dir}/{repo}/test.yaml") as yaml_test:
                values = load_hyperpyyaml(yaml_test)

            print(f"Collecting results for: {repo} w/ values={values}")
            prediction = get_prediction(repo, values)

            # extend the results
            results[repo] = {"before": prediction}
            with open(yaml_path, "w") as yaml_out:
                yaml.dump(results, yaml_out, default_flow_style=None)


def gather_refactoring_results(
    glob_filter="*",
    new_interfaces_git="https://github.com/speechbrain/speechbrain",
    new_interfaces_branch="hf-interface-testing",
    new_interfaces_local_dir="tests/tmp/hf_interfaces",
    yaml_path="tests/tmp/refactoring_results.yaml",
):
    """After refactoring HF YAMLs and/or code, gather prediction results.

    Parameters
    ----------
    glob_filter: str
        Filter for a repo subset or a specific repo.
    new_interfaces_git: str
        Your git repo (or default: `https://github.com/speechbrain/speechbrain`);
        can be specified in tests/utils/overrides.yaml
    new_interfaces_branch: str
        Default is `hf-interface-testing` (a git branch); can be specified in tests/utils/overrides.yaml
    new_interfaces_local_dir: str
        Default is `tests/tmp/hf_interfaces` (a local path); can be specified in tests/utils/overrides.yaml
    yaml_path: str
        Path where to store/load refactoring testing results for later comparison.
    """
    # expected results need to exist
    if os.path.exists(yaml_path):
        with open(yaml_path) as yaml_in:
            results = yaml.safe_load(yaml_in)

    # go through each repo
    updates_dir = init(
        new_interfaces_git, new_interfaces_branch, new_interfaces_local_dir
    )
    repos = map(os.path.basename, glob(f"{updates_dir}/{glob_filter}"),)
    for repo in repos:
        # skip if results are there
        if "after" not in results[repo].keys():
            # get values
            with open(f"{updates_dir}/{repo}/test.yaml") as yaml_test:
                values = load_hyperpyyaml(yaml_test)

            print(
                f"Collecting refactoring results for: {repo} w/ values={values}"
            )

            # extend the results
            results[repo]["after"] = get_prediction(repo, values, updates_dir)
            results[repo]["same"] = (
                results[repo]["before"] == results[repo]["after"]
            )

            # update
            with open(yaml_path, "w") as yaml_out:
                yaml.dump(results, yaml_out, default_flow_style=None)

            print(f"\tsame: {results[repo]['same'] }")


def test_performance(
    repo, values, run_opts, updates_dir=None, recipe_overrides={}
):
    """Runs the evaluation partition of a recipe dataset for a pretrained model.

    Parameters
    ----------
    repo: str
        Source of pretrained model (assuming its within the HF speechbrain collection).
    values: dict
        Interface specification.
        Examples: speechbrain:hf-interface-testing/updates_pretrained_models/ssl-wav2vec2-base-librispeech/test.yaml
                  speechbrain:hf-interface-testing/updates_pretrained_models/asr-wav2vec2-librispeech/test.yaml
    run_opts: dict
        Run options, such as device
    updates_dir: str
        Controls whether/not we are in the refactored results (None: expected results; before refactoring).
    recipe_overrides: dict
        Recipe YAMLs contain placeholders and flags which need to be overwritten (e.g. data_folder & skip_prep).
        See: overrides.yaml

    Returns
    -------
    Dict for export to yaml with performance statistics, as specified in the test.yaml files.
    """
    # Dataset depending file structure
    tmp_dir = f'tests/tmp/{values["dataset"]}'
    speechbrain.create_experiment_directory(experiment_directory=tmp_dir)
    stats_meta = {
        f'[{values["dataset"]}] - {"BEFORE" if updates_dir is None else "AFTER"}': repo
    }

    # Load pretrained
    model = get_model(
        repo=repo, values=values, updates_dir=updates_dir, run_opts=run_opts
    )  # noqa

    # Dataio preparation; we need the test sets only
    with open(values["recipe_yaml"]) as fin:
        recipe_hparams = load_hyperpyyaml(
            fin, values["overrides"] | recipe_overrides
        )

    # Dataset preparation is assumed to be done through recipes; before running this.
    exec(values["dataio"])
    test_datasets = deepcopy(eval(values["test_datasets"]))

    # harmonise
    if type(test_datasets) is not dict:
        tmp = {}
        if type(test_datasets) is list:
            for i, x in enumerate(test_datasets):
                tmp[i] = x
        else:
            tmp[0] = test_datasets
        test_datasets = tmp

    # prepare testing
    logger = FileTrainLogger(save_file=f"{tmp_dir}/{repo}.log")
    reporting = deepcopy(values["performance"])
    for metric, specs in reporting.items():
        reporting[metric]["tracker"] = deepcopy(
            recipe_hparams[specs["handler"]]()
        )

    test_loader_kwargs = deepcopy(recipe_hparams[values["test_loader"]])
    del recipe_hparams

    stats = {}
    for k in test_datasets.keys():  # keys are test_clean, test_other etc
        test_set = test_datasets[k]
        if not (
            isinstance(test_set, DataLoader)
            or isinstance(test_set, LoopedLoader)
        ):
            test_set = make_dataloader(test_set, **test_loader_kwargs)

        with torch.no_grad():
            for batch in tqdm(test_set, dynamic_ncols=True, disable=False):
                batch = batch.to(model.device)
                wavs, wav_lens = batch.sig
                wavs, wav_lens = (  # noqa
                    wavs.to(model.device),
                    wav_lens.to(model.device),
                )
                predictions = eval(  # noqa
                    f'model.{values["fnx"]}(wavs, wav_lens)'
                )
                predicted = eval(values["predicted"])  # noqa
                targeted = eval(values["targeted"])  # noqa
                ids = batch.id  # noqa
                for metric in reporting.keys():
                    reporting[metric]["tracker"].append(
                        *eval(values["to_stats"])
                    )

        stats[k] = {}
        for metric, specs in reporting.items():
            stats[k][metric] = specs["tracker"].summarize(specs["field"])
        logger.log_stats(
            stats_meta=stats_meta | {"set": k}, test_stats=stats[k],
        )

    return stats


# run first w/ "--after=False" on latest develop, then checkout the refactoring branch and run w/ "--after=True"
# PYTHONPATH=`realpath .` python tests/utils/refactoring_checks.py tests/utils/overrides.yaml --LibriSpeech_data="" --CommonVoice_EN_data="" --CommonVoice_FR_data="" --IEMOCAP_data="" --after=False
if __name__ == "__main__":
    hparams_file, run_opts, overrides = speechbrain.parse_arguments(
        sys.argv[1:]
    )

    with open(hparams_file) as fin:
        dataset_overrides = load_hyperpyyaml(fin, overrides)

    # go through each repo
    updates_dir = init(
        dataset_overrides["new_interfaces_git"],
        dataset_overrides["new_interfaces_branch"],
        dataset_overrides["new_interfaces_local_dir"],
    )

    # load results, if existing -or- new from scratch
    yaml_path = f'{dataset_overrides["new_interfaces_local_dir"]}.yaml'
    if os.path.exists(yaml_path):
        with open(yaml_path) as yaml_in:
            results = yaml.safe_load(yaml_in)
    else:
        results = {}

    repos = map(
        os.path.basename,
        glob(f'{updates_dir}/{dataset_overrides["glob_filter"]}'),
    )
    for repo in repos:
        # get values
        with open(f"{updates_dir}/{repo}/test.yaml") as yaml_test:
            values = load_hyperpyyaml(yaml_test)

        # for this testing, some fields need to exist; skip otherwise
        if any(
            [
                entry not in values
                for entry in [
                    "dataset",
                    "overrides",
                    "dataio",
                    "test_datasets",
                    "test_loader",
                    "performance",
                    "predicted",
                ]
            ]
        ):
            continue

        # skip if datasets is not given
        if not dataset_overrides[f'{values["dataset"]}_data']:
            continue

        print(f"Run tests on: {repo}")
        if repo not in results.keys():
            results[repo] = {}

        # Before refactoring
        if "before" not in results[repo].keys():
            results[repo]["before"] = test_performance(
                repo,
                values,
                updates_dir=None,
                run_opts=run_opts,
                recipe_overrides=dataset_overrides[values["dataset"]],
            )

            # update
            with open(yaml_path, "w") as yaml_out:
                yaml.dump(results, yaml_out, default_flow_style=None)

        # After refactoring
        if (
            "after" not in results[repo].keys()
            and dataset_overrides["after"] is True
        ):
            results[repo]["after"] = test_performance(
                repo,
                values,
                run_opts=run_opts,
                updates_dir=updates_dir,
                recipe_overrides=dataset_overrides[values["dataset"]],
            )

            results[repo]["same"] = (
                results[repo]["before"] == results[repo]["after"]
            )
            print(f'\tbefore: {results[repo]["before"]}')
            print(f'\t after: {results[repo]["after"]}')
            print(f'\t  same: {results[repo]["same"]}')

            # update
            with open(yaml_path, "w") as yaml_out:
                yaml.dump(results, yaml_out, default_flow_style=None)

    # update
    with open(yaml_path, "w") as yaml_out:
        yaml.dump(results, yaml_out, default_flow_style=None)

tests/utils/refactoring_checks.py
TzuchengChang/NASS