wehub/nvidia-nemo--speech
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
Details
CICD NeMo / pre-flight (push) Failing after 0s
Details
CICD NeMo / configure (push) Has been skipped
Details
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
Details
CICD NeMo / code-linting (push) Has been skipped
Details
Build, validate, and release Neural Modules / release (push) Has been skipped
Details
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
Details
CICD NeMo / cicd-test-container-build (push) Has been cancelled
Details
CICD NeMo / cicd-import-tests (push) Has been cancelled
Details
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
Details
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
Details
CICD NeMo / cicd-main-speech (push) Has been cancelled
Details
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
Details
CICD NeMo / Coverage (e2e) (push) Has been cancelled
Details
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
Details
CodeQL / Analyze (python) (push) Has been cancelled
Details
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:58 +08:00
commit ba4be087d5
2316 changed files with 2668701 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tutorials/audio/README.md
+6
View File
@@ -0,0 +1,6 @@
# NeMo Audio Processing Tutorials
An quick way to get started with audio processing in NeMo.
## Speech Enhancement
* [Quickstart: Speech Enhancement with NeMo](speech_enhancement/Speech_Enhancement_with_NeMo.ipynb)
tutorials/audio/speech_enhancement/BNR_Speech_enhancement_with_NeMo.ipynb
+880
View File
@@ -0,0 +1,880 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "2fb43f21",
"metadata": {
"id": "2fb43f21"
},
"source": [
"# Introduction\n",
"\n",
"The goal of this tutorial is to demonstrate the basic steps required to setup and train the Maxine Background Noise removal network [1] in NeMo.\n",
"\n",
"This notebook covers the following steps:\n",
"\n",
"* Download speech and noise data\n",
"* Prepare the training data by mixing speech and noise\n",
"* Configure and train a simple single-output model\n",
"\n",
"Note that this tutorial is only for demonstration purposes.\n",
"\n",
"*Disclaimer:*\n",
"User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "50525ece",
"metadata": {
"id": "50525ece"
},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
"\n",
"Instructions for setting up Colab are as follows:\n",
"1. Open a new Python 3 notebook.\n",
"2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
"3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
"4. Run this cell to set up dependencies.\n",
"5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
"\"\"\"\n",
"\n",
"GIT_USER = 'NVIDIA-NeMo'\n",
"BRANCH = 'main'\n",
"\n",
"if 'google.colab' in str(get_ipython()):\n",
"\n",
" # Install dependencies\n",
" !pip install wget\n",
" !apt-get install sox libsndfile1 ffmpeg\n",
" !pip install text-unidecode\n",
" !pip install matplotlib>=3.3.2\n",
"\n",
" ## Install NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/{GIT_USER}/Speech.git@{BRANCH}\"\n",
"\n",
" ## Install TorchAudio\n",
" !pip install torchaudio>=0.13.0 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "765e8b8a",
"metadata": {
"id": "765e8b8a"
},
"source": [
"The following cell will take care of the necessary imports and prepare utility functions used throughout the notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9248dd5f",
"metadata": {
"id": "9248dd5f"
},
"outputs": [],
"source": [
"import glob\n",
"import librosa\n",
"import os\n",
"import torch\n",
"import tqdm\n",
"\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import lightning.pytorch as pl\n",
"import soundfile as sf\n",
"\n",
"from omegaconf import OmegaConf, open_dict\n",
"from pathlib import Path\n",
"from torchmetrics.functional.audio import signal_distortion_ratio, scale_invariant_signal_distortion_ratio\n",
"\n",
"from nemo.utils.notebook_utils import download_an4\n",
"from nemo.collections.asr.parts.preprocessing.segment import AudioSegment\n",
"from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest\n",
"\n",
"\n",
"# Utility functions for displaying signals and metrics\n",
"def show_signal(signal: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show the time-domain signal and its spectrogram.\n",
" \"\"\"\n",
" fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 2.5))\n",
"\n",
" # show waveform\n",
" t = np.arange(0, len(signal)) / sample_rate\n",
"\n",
" ax[0].plot(t, signal)\n",
" ax[0].set_xlim(0, t.max())\n",
" ax[0].grid()\n",
" ax[0].set_xlabel('time / s')\n",
" ax[0].set_ylabel('amplitude')\n",
" ax[0].set_title(tag)\n",
"\n",
" n_fft = 1024\n",
" hop_length = 256\n",
"\n",
" D = librosa.amplitude_to_db(np.abs(librosa.stft(signal, n_fft=n_fft, hop_length=hop_length)), ref=np.max)\n",
" img = librosa.display.specshow(D, y_axis='linear', x_axis='time', sr=sample_rate, n_fft=n_fft, hop_length=hop_length, ax=ax[1])\n",
" ax[1].set_title(tag)\n",
" \n",
" plt.tight_layout()\n",
" plt.colorbar(img, format=\"%+2.f dB\", ax=ax)\n",
"\n",
"def show_metrics(signal: np.ndarray, reference: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show metrics for the time-domain signal and the reference signal.\n",
" \"\"\"\n",
" sdr = signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" sisdr = scale_invariant_signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" print(tag)\n",
" print('\\tsdr: ', sdr.item())\n",
" print('\\tsisdr:', sisdr.item())"
]
},
{
"cell_type": "markdown",
"id": "bfa2199e",
"metadata": {
"id": "bfa2199e"
},
"source": [
"### Data preparation"
]
},
{
"cell_type": "markdown",
"id": "0aa05a36",
"metadata": {
"id": "0aa05a36"
},
"source": [
"In this notebook, it is assumed that all audio will be resampled to 16kHz and the data and configuration will be stored under `root_dir` as defined below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8e790630",
"metadata": {
"id": "8e790630"
},
"outputs": [],
"source": [
"# sample rate used throughout the notebook\n",
"sample_rate = 16000\n",
"\n",
"# root directory for data preparation, configurations, etc\n",
"root_dir = Path('./')\n",
"\n",
"# data directory\n",
"data_dir = root_dir / 'data'\n",
"data_dir.mkdir(exist_ok=True)\n",
"\n",
"# scripts directory\n",
"scripts_dir = root_dir / 'scripts'\n",
"scripts_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "markdown",
"id": "9af717ad",
"metadata": {
"id": "9af717ad"
},
"source": [
"Clean speech data is used to prepare datasets used for training a simple speech enhancement model.\n",
"\n",
"In this tutorial, a subset of LibriSpeech dataset [2] will be downloaded and used as the speech material.\n",
"The following cell will download and prepare the speech data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1c5a60a",
"metadata": {
"id": "e1c5a60a"
},
"outputs": [],
"source": [
"speech_dir = data_dir / 'speech'\n",
"speech_data_set = 'mini'\n",
"\n",
"# Copy script\n",
"get_librispeech_script = os.path.join(scripts_dir, 'get_librispeech_data.py')\n",
"if not os.path.exists(get_librispeech_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_librispeech_data.py\n",
"\n",
"# Download the data\n",
"if not speech_dir.is_dir():\n",
" speech_dir.mkdir(exist_ok=True)\n",
" !python {get_librispeech_script} --data_root={speech_dir} --data_set={speech_data_set}\n",
"else:\n",
" print('Speech dataset already exists in:', speech_dir)\n",
"\n",
"# Reduce the size of test dataset for this tutorial to 1000 clean utterances for train and 100 clean utterances for test\n",
"train_metadata = read_manifest(speech_dir / 'train_clean_5.json')\n",
"write_manifest(speech_dir / 'train.json', train_metadata[:1000])\n",
"\n",
"test_metadata = read_manifest(speech_dir / 'dev_clean_2.json')\n",
"write_manifest(speech_dir / 'test.json', test_metadata[:100])\n",
"\n",
"# Speech manifests\n",
"speech_manifest = {\n",
" 'train': speech_dir / 'train.json',\n",
" 'test': speech_dir / 'test.json',\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "06f07374",
"metadata": {
"id": "06f07374"
},
"source": [
"Noise data will be mixed with the downloaded speech data to prepare a noisy dataset.\n",
"\n",
"The following cell will download and prepare the noise data using a subset of the DEMAND dataset [3] will be downloaded and used as the noise data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1de1089",
"metadata": {
"id": "b1de1089"
},
"outputs": [],
"source": [
"noise_dir = data_dir / 'noise'\n",
"noise_data_set = 'STRAFFIC,PSTATION'\n",
"\n",
"# Copy script\n",
"get_demand_script = os.path.join(scripts_dir, 'get_demand_data.py')\n",
"if not os.path.exists(get_demand_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_demand_data.py\n",
"\n",
"if not noise_dir.is_dir():\n",
" noise_dir.mkdir(exist_ok=True)\n",
" !python {get_demand_script} --data_root={noise_dir} --data_sets={noise_data_set}\n",
"else:\n",
" print('Noise directory already exists in:', noise_dir)\n",
"\n",
"\n",
"def create_noise_manifest(base_dir, subset, offset=0, duration=None):\n",
" \"\"\"Split the noise data set into train and test subsets.\n",
" \"\"\"\n",
" complete_noise_manifests = glob.glob(str(base_dir / 'manifests' / '*.json'))\n",
" subset_noise_manifest = base_dir / f'{subset}_manifest.json'\n",
" \n",
" subset_metadata = []\n",
"\n",
" for noise_manifest in complete_noise_manifests:\n",
" complete_metadata = read_manifest(noise_manifest)\n",
" \n",
" for item in complete_metadata:\n",
" new_item = item.copy()\n",
" new_item['offset'] = offset\n",
" new_item['duration'] = duration\n",
" subset_metadata.append(new_item)\n",
"\n",
" write_manifest(subset_noise_manifest.as_posix(), subset_metadata)\n",
"\n",
" return subset_noise_manifest\n",
"\n",
"noise_manifest = {\n",
" 'train': create_noise_manifest(noise_dir, 'train', offset=0, duration=200),\n",
" 'test': create_noise_manifest(noise_dir, 'test', offset=200, duration=100),\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "1378d78a",
"metadata": {
"id": "1378d78a"
},
"source": [
"For this tutorial, a single-channel noisy dataset is constructed by adding speech and noise.\n",
"\n",
"The following block will add speech and noise and save the noisy data. The noisy data is created by mixing speech and noise at a few pre-defined signal-to-noise ratios (SNRs). Note that a separate manifest will be created for each SNR."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d6c678a0",
"metadata": {
"id": "d6c678a0"
},
"outputs": [],
"source": [
"%%capture\n",
"# Suppress output of this cell, since the script used below is relatively verbose.\n",
"\n",
"# Copy script\n",
"add_noise_script = os.path.join(scripts_dir, 'add_noise.py')\n",
"if not os.path.exists(add_noise_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/add_noise.py\n",
"\n",
"# Generate noisy datasets and save the noise component as well.\n",
"noisy_dir = data_dir / 'noisy'\n",
"noisy_dir.mkdir(exist_ok=True)\n",
"\n",
"for subset in ['train', 'test']:\n",
" noisy_subset_dir = noisy_dir / subset\n",
"\n",
" if not noisy_subset_dir.is_dir():\n",
" noisy_subset_dir.mkdir(exist_ok=True)\n",
" !python {add_noise_script} --input_manifest={speech_manifest[subset]} --noise_manifest={noise_manifest[subset]} --out_dir={noisy_subset_dir} --snrs 0 5 10 15 20 --num_workers 4 --save_noise"
]
},
{
"cell_type": "markdown",
"id": "c4cd1426",
"metadata": {
"id": "c4cd1426"
},
"source": [
"Training a model requires an input dataset which includes information about the noisy input signal and the desired (target) output signal.\n",
"\n",
"In this tutorial, train and test manifests are created by combining the information from the speech manifests and each noisy manifest generated in the previous step. Note that the final manifests include `noisy_filepath`, `speech_filepath` and `noise_filepath`. These keys can be used to define the input signal and the output signal for the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fcfeb639",
"metadata": {
"id": "fcfeb639"
},
"outputs": [],
"source": [
"dataset_manifest = {\n",
" 'train': data_dir / 'dataset_train.json',\n",
" 'test': data_dir / 'dataset_test.json',\n",
"}\n",
"\n",
"for subset in ['train', 'test']:\n",
" # Load clean manifest\n",
" speech_metadata = read_manifest(speech_manifest[subset])\n",
"\n",
" # Load noisy manifests\n",
" noisy_manifests = glob.glob(str(noisy_dir / subset / 'manifests/*.json'))\n",
" noisy_manifests.sort()\n",
"\n",
" subset_metadata = []\n",
"\n",
" for noisy_manifest in noisy_manifests:\n",
" noisy_metadata = read_manifest(noisy_manifest)\n",
"\n",
" for speech_item, noisy_item in tqdm.tqdm(zip(speech_metadata, noisy_metadata), total=len(noisy_metadata)):\n",
" # Check that the file matches\n",
" assert os.path.basename(speech_item['audio_filepath']) == os.path.basename(noisy_item['audio_filepath']), f'Speech: {speech_item}. Noisy: {noisy_item}'\n",
"\n",
" # Create a new item for the subset manifest\n",
" subset_item = {\n",
" 'noisy_filepath': noisy_item['audio_filepath'],\n",
" 'speech_filepath': speech_item['audio_filepath'],\n",
" 'noise_filepath': noisy_item['noise_filepath'],\n",
" 'duration': noisy_item['duration'],\n",
" 'offset': noisy_item.get('offset', 0)\n",
" }\n",
"\n",
" subset_metadata.append(subset_item)\n",
"\n",
" # Save the subset manifest\n",
" write_manifest(dataset_manifest[subset].as_posix(), subset_metadata)"
]
},
{
"cell_type": "markdown",
"id": "aa4248ab",
"metadata": {
"id": "aa4248ab"
},
"source": [
"### Model configuration\n",
"\n",
"We use the SEASR model architecture for training the dataset.",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "002b917c",
"metadata": {
"id": "002b917c"
},
"outputs": [],
"source": [
"config_dir = root_dir / 'conf'\n",
"config_dir.mkdir(exist_ok=True)\n",
"\n",
"config_path = config_dir / 'maxine_bnr.yaml'\n",
"\n",
"if not config_path.is_file():\n",
" !wget https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/examples/audio/conf/maxine_bnr.yaml -P {config_dir.as_posix()}\n",
"\n",
"config = OmegaConf.load(config_path)\n",
"config = OmegaConf.to_container(config, resolve=True)\n",
"config = OmegaConf.create(config)\n",
"\n",
"print('Loaded config')\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"id": "65cab3ad",
"metadata": {
"id": "65cab3ad"
},
"source": [
"Training dataset is configured with the following parameters\n",
"* `manifest_filepath` points to a manifest file, with each line containing a dictionary corresponding to a single example\n",
"* `input_key` is the key corresponding to the input audio signal in the example dictionary\n",
"* `target_key` is the key corresponding to the desired output (target) audio signal in the example dictionary\n",
"* `min_duration` can be used to filter out short examples"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "91717b15",
"metadata": {
"id": "91717b15"
},
"outputs": [],
"source": [
"# Setup training dataset\n",
"config.model.train_ds.manifest_filepath = dataset_manifest['train'].as_posix()\n",
"config.model.train_ds.input_key = 'noisy_filepath'\n",
"config.model.train_ds.target_key = 'speech_filepath'\n",
"config.model.train_ds.min_duration = 0 # load all audio files, without filtering short ones\n",
"\n",
"print(\"Train dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.train_ds))"
]
},
{
"cell_type": "markdown",
"id": "c11592e4",
"metadata": {
"id": "c11592e4"
},
"source": [
"Validation and test datasets can be configured in the same way as the training dataset. Here, we use the same dataset for validation and testing purposes for simplicity."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d8bc0a0a",
"metadata": {
"id": "d8bc0a0a"
},
"outputs": [],
"source": [
"# Use test manifest for validation and test sets\n",
"config.model.validation_ds.manifest_filepath = dataset_manifest['test'].as_posix()\n",
"config.model.validation_ds.input_key = 'noisy_filepath'\n",
"config.model.validation_ds.target_key = 'speech_filepath'\n",
"\n",
"config.model.test_ds.manifest_filepath = dataset_manifest['test'].as_posix()\n",
"config.model.test_ds.input_key = 'noisy_filepath'\n",
"config.model.test_ds.target_key = 'speech_filepath'\n",
"\n",
"print(\"Validation dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.validation_ds))\n",
"\n",
"print(\"Test dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.test_ds))"
]
},
{
"cell_type": "markdown",
"id": "1bc0f8f3",
"metadata": {
"id": "1bc0f8f3"
},
"source": [
"Metrics for validation and test set are configured in the following cell.\n",
"\n",
"In this tutorial, signal-to-distortion ratio (SDR) and scale-invariant SDR from torch metrics are used [4]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "82de1dbb",
"metadata": {
"id": "82de1dbb"
},
"outputs": [],
"source": [
"# Setup metrics to compute on validation and test sets\n",
"metrics = OmegaConf.create({\n",
" 'sisdr': {\n",
" '_target_': 'torchmetrics.audio.ScaleInvariantSignalDistortionRatio',\n",
" },\n",
" 'sdr': {\n",
" '_target_': 'torchmetrics.audio.SignalDistortionRatio',\n",
" }\n",
"})\n",
"config.model.metrics.val = metrics\n",
"config.model.metrics.test = metrics\n",
"\n",
"print(\"Metrics config:\")\n",
"print(OmegaConf.to_yaml(config.model.metrics))"
]
},
{
"cell_type": "markdown",
"id": "4eda2d65",
"metadata": {
"id": "4eda2d65"
},
"source": [
"### Trainer configuration\n",
"NeMo models are primarily PyTorch Lightning modules and therefore are entirely compatible with the PyTorch Lightning ecosystem."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e53ba6ef",
"metadata": {
"id": "e53ba6ef"
},
"outputs": [],
"source": [
"print(\"Trainer config:\")\n",
"print(OmegaConf.to_yaml(config.trainer))"
]
},
{
"cell_type": "markdown",
"id": "a9201689",
"metadata": {
"id": "a9201689"
},
"source": [
"We can modify some trainer configs for this tutorial.\n",
"Most importantly, the number of epochs is set to a small value, to limit the runtime for the purpose of this example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "17b7499f",
"metadata": {
"id": "17b7499f"
},
"outputs": [],
"source": [
"# Checks if we have GPU available and uses it\n",
"accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
"config.trainer.devices = 1\n",
"config.trainer.accelerator = accelerator\n",
"\n",
"# Reduces maximum number of epochs for quick demonstration\n",
"config.trainer.max_epochs = 10\n",
"\n",
"# Remove distributed training flags\n",
"config.trainer.strategy = 'auto'\n",
"\n",
"# Instantiate the trainer\n",
"trainer = pl.Trainer(**config.trainer)"
]
},
{
"cell_type": "markdown",
"id": "053192a4",
"metadata": {
"id": "053192a4"
},
"source": [
"### Experiment manager\n",
"\n",
"NeMo has an experiment manager that handles logging and checkpointing."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ccb8948d",
"metadata": {
"id": "ccb8948d"
},
"outputs": [],
"source": [
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"# The exp_dir provides a path to the current experiment for easy access\n",
"\n",
"print(\"Experiment directory:\")\n",
"print(exp_dir)"
]
},
{
"cell_type": "markdown",
"id": "8bb583f8",
"metadata": {
"id": "8bb583f8"
},
"source": [
"### Model instantiation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6905fea2",
"metadata": {
"id": "6905fea2"
},
"outputs": [],
"source": [
"from nemo.collections.audio.models.maxine import BNR2",
"\n",
"enhancement_model = BNR2(cfg=config.model, trainer=trainer)"
]
},
{
"cell_type": "markdown",
"id": "6ee1a6d2",
"metadata": {
"id": "6ee1a6d2"
},
"source": [
"### Training\n",
"Create a Tensorboard visualization to monitor progress"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dc9e7639",
"metadata": {
"id": "dc9e7639"
},
"outputs": [],
"source": [
"try:\n",
" from google import colab\n",
" COLAB_ENV = True\n",
"except (ImportError, ModuleNotFoundError):\n",
" COLAB_ENV = False\n",
"\n",
"# Load the TensorBoard notebook extension\n",
"if COLAB_ENV:\n",
" %load_ext tensorboard\n",
" %tensorboard --logdir {exp_dir}\n",
"else:\n",
" print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
]
},
{
"cell_type": "markdown",
"id": "5e962b9b",
"metadata": {
"id": "5e962b9b"
},
"source": [
"Training can be started using `trainer.fit`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b00b3f4",
"metadata": {
"id": "0b00b3f4"
},
"outputs": [],
"source": [
"trainer.fit(enhancement_model)"
]
},
{
"cell_type": "markdown",
"id": "5b5c1e78",
"metadata": {
"id": "5b5c1e78"
},
"source": [
"After the training is completed, the configured metrics can be easily computed on the test set as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5216195d",
"metadata": {
"id": "5216195d"
},
"outputs": [],
"source": [
"trainer.test(enhancement_model, ckpt_path=None)"
]
},
{
"cell_type": "markdown",
"id": "bf5b73e6",
"metadata": {
"id": "bf5b73e6"
},
"source": [
"### Inference\n",
"\n",
"The following cell provides an example of inference on an single audio file.\n",
"For simplicity, the audio file information is taken from the test dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b5cf923",
"metadata": {
"id": "6b5cf923"
},
"outputs": [],
"source": [
"# Load a single audio example from the test set\n",
"test_metadata = read_manifest(dataset_manifest['test'].as_posix())\n",
"\n",
"# Path to audio files\n",
"noisy_filepath = test_metadata[-1]['noisy_filepath'] # noisy audio\n",
"speech_filepath = test_metadata[-1]['speech_filepath'] # clean speech\n",
"noise_filepath = test_metadata[-1]['noise_filepath'] # corresponding noise\n",
"\n",
"# Load audio\n",
"noisy_signal = AudioSegment.from_file(noisy_filepath, target_sr=sample_rate).samples\n",
"speech_signal = AudioSegment.from_file(speech_filepath, target_sr=sample_rate).samples\n",
"\n",
"# Move to device\n",
"device = 'cuda' if accelerator == 'gpu' else 'cpu'\n",
"enhancement_model = enhancement_model.to(device)\n",
"\n",
"# Process using the model\n",
"noisy_tensor = torch.tensor(noisy_signal).reshape(1, 1, -1).to(device) # (batch, channel, time)\n",
"with torch.no_grad():\n",
" output_tensor = enhancement_model(input_signal=noisy_tensor)\n",
"output_signal = output_tensor[0][0].detach().cpu().numpy()"
]
},
{
"cell_type": "markdown",
"id": "90f6c0d0",
"metadata": {
"id": "90f6c0d0"
},
"source": [
"Signals can be easily plotted and signal metrics can be calculated for the given example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4aec5c64",
"metadata": {
"id": "4aec5c64"
},
"outputs": [],
"source": [
"# Show noisy and clean signals\n",
"show_metrics(signal=noisy_signal, reference=speech_signal, tag='Noisy signal', sample_rate=sample_rate)\n",
"show_metrics(signal=output_signal, reference=speech_signal, tag='Output signal', sample_rate=sample_rate)\n",
"\n",
"# Show signals\n",
"show_signal(speech_signal, tag='Speech signal')\n",
"show_signal(noisy_signal, tag='Noisy signal')\n",
"show_signal(output_signal, tag='Output signal')\n",
"\n",
"# Play audio\n",
"print('Speech signal')\n",
"ipd.display(ipd.Audio(speech_signal, rate=sample_rate))\n",
"\n",
"print('Noisy signal')\n",
"ipd.display(ipd.Audio(noisy_signal, rate=sample_rate))\n",
"\n",
"print('Output signal')\n",
"ipd.display(ipd.Audio(output_signal, rate=sample_rate))"
]
},
{
"cell_type": "markdown",
"id": "cd6da1f0",
"metadata": {
"id": "cd6da1f0"
},
"source": [
"## Next steps\n",
"This is a simple tutorial which can serve as a starting point for prototyping and experimentation with audio-to-audio models.\n",
"A processed audio output can be used, for example, for ASR or TTS.\n",
"\n",
"For more details about NeMo models and applications in in ASR and TTS, we recommend you checkout other tutorials next:\n",
"\n",
"* [NeMo fundamentals](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb)\n",
"* [NeMo models](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb)\n",
"* [Speech Recognition](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)\n",
"* [Speech Synthesis](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/tts/Inference_ModelSelect.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "46a855e3",
"metadata": {
"id": "46a855e3"
},
"source": [
"## References\n",
"\n",
"[1] M. Remane, R. R. Nalla and A. Dantrey, \"SEASR: Speech Enhancement for Automatic Speech Recognition Systems using Convolution Recurrent Neural Network with Residual Connections,\" 2024 IEEE 5th Women in Technology Conference (WINTECHCON), Bengaluru, India, 2024, pp. 1-5, doi: 10.1109/Wintechcon61988.2024.10837982\n",
"\n",
"[2] V. Panayotov, G. Chen, D. Povery, S. Khudanpur, \"LibriSpeech: An ASR corpus based on public domain audio books,\" ICASSP 2015\n",
"\n",
"[3] J. Thieman, N. Ito, V. Emmanuel, \"DEMAND: collection of multi-channel recordings of acoustic noise in diverse environments,\" ICA 2013\n",
"\n",
"[4] https://github.com/Lightning-AI/torchmetrics\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"gpuClass": "standard",
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/audio/speech_enhancement/Speech_Enhancement_with_NeMo.ipynb
+1318
View File
File diff suppressed because it is too large Load Diff
tutorials/audio/speech_enhancement/Speech_Enhancement_with_Online_Augmentation.ipynb
+984
View File
@@ -0,0 +1,984 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "3b53ae0a-4cb5-44fd-bcde-5ce7abce79e9",
"metadata": {
"id": "3b53ae0a-4cb5-44fd-bcde-5ce7abce79e9"
},
"source": [
"# Introduction\n",
"\n",
"The goal of this tutorial is to demonstrate the basic steps required to setup and train a simple single-channel speech enhancement model in NeMo using online augmentation with noise and room impulse responce (RIR). Online augmentation is performed using a dataloader based on Lhotse speech data processing toolkit [1].\n",
"\n",
"\n",
"This notebook covers the following steps:\n",
"\n",
"* Download speech, noise and RIR data\n",
"* Prepare Lhotse manifests for speech, noise and RIR data\n",
"* Prepare fixed validation set by mixing speech, noise and RIR data\n",
"* Configure and train a simple single-output model\n",
"\n",
"Note that this tutorial is only for demonstration purposes.\n",
"To achieve best performance for a particular use case, carefully prepared data and more advanced models should be used.\n",
"\n",
"*Disclaimer:*\n",
"User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d603b38-b63c-42a6-92b3-575925149bde",
"metadata": {
"id": "3d603b38-b63c-42a6-92b3-575925149bde"
},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
"\n",
"Instructions for setting up Colab are as follows:\n",
"1. Open a new Python 3 notebook.\n",
"2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
"3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
"4. Run this cell to set up dependencies.\n",
"5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
"\"\"\"\n",
"\n",
"GIT_USER = 'NVIDIA-NeMo'\n",
"BRANCH = 'main'\n",
"\n",
"if 'google.colab' in str(get_ipython()):\n",
"\n",
" # Install dependencies\n",
" !pip install wget\n",
" !apt-get install sox libsndfile1 ffmpeg\n",
" !pip install text-unidecode\n",
" !pip install matplotlib>=3.3.2\n",
"\n",
" ## Install NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/{GIT_USER}/Speech.git@{BRANCH}\"\n",
"\n",
" ## Install TorchAudio\n",
" !pip install torchaudio>=0.13.0 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "bdac6ac8-a21f-4ea8-8484-41a6d187a2fe",
"metadata": {
"id": "bdac6ac8-a21f-4ea8-8484-41a6d187a2fe"
},
"source": [
"The following cell will take care of the necessary imports and prepare utility functions used throughout the notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "baca13c6-b0ed-429b-93b6-77249fdf4710",
"metadata": {
"id": "baca13c6-b0ed-429b-93b6-77249fdf4710"
},
"outputs": [],
"source": [
"import glob\n",
"import librosa\n",
"import os\n",
"import torch\n",
"import tqdm\n",
"from itertools import islice\n",
"\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import lightning.pytorch as pl\n",
"import soundfile as sf\n",
"from pathlib import Path\n",
"from omegaconf import OmegaConf, open_dict\n",
"from sklearn.model_selection import train_test_split\n",
"from torchmetrics.functional.audio import signal_distortion_ratio, scale_invariant_signal_distortion_ratio\n",
"from lhotse import CutSet, RecordingSet, Recording, MonoCut\n",
"from lhotse.recipes import (\n",
" download_rir_noise,\n",
" prepare_rir_noise,\n",
" download_librispeech,\n",
" prepare_librispeech\n",
")\n",
"\n",
"from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config\n",
"from nemo.collections.audio.data.audio_to_audio_lhotse import LhotseAudioToTargetDataset"
]
},
{
"cell_type": "markdown",
"id": "fc0fafa4-bc65-4066-8e75-d740e2d15259",
"metadata": {
"id": "fc0fafa4-bc65-4066-8e75-d740e2d15259"
},
"source": [
"Utility functions for displaying signals and metrics"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "49720c06-3b4d-45b6-a054-73800c3bddea",
"metadata": {
"id": "49720c06-3b4d-45b6-a054-73800c3bddea"
},
"outputs": [],
"source": [
"def show_signal(signal: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show the time-domain signal and its spectrogram.\n",
" \"\"\"\n",
" fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 2.5))\n",
"\n",
" # show waveform\n",
" t = np.arange(0, len(signal)) / sample_rate\n",
"\n",
" ax[0].plot(t, signal)\n",
" ax[0].set_xlim(0, t.max())\n",
" ax[0].grid()\n",
" ax[0].set_xlabel('time / s')\n",
" ax[0].set_ylabel('amplitude')\n",
" ax[0].set_title(tag)\n",
"\n",
" n_fft = 1024\n",
" hop_length = 256\n",
"\n",
" D = librosa.amplitude_to_db(np.abs(librosa.stft(signal, n_fft=n_fft, hop_length=hop_length)), ref=np.max)\n",
" img = librosa.display.specshow(D, y_axis='linear', x_axis='time', sr=sample_rate, n_fft=n_fft, hop_length=hop_length, ax=ax[1])\n",
" ax[1].set_title(tag)\n",
"\n",
" plt.tight_layout()\n",
" plt.colorbar(img, format=\"%+2.f dB\", ax=ax)\n",
"\n",
"def show_metrics(signal: np.ndarray, reference: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show metrics for the time-domain signal and the reference signal.\n",
" \"\"\"\n",
" sdr = signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" sisdr = scale_invariant_signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" print(tag)\n",
" print('\\tsdr: ', sdr.item())\n",
" print('\\tsisdr:', sisdr.item())"
]
},
{
"cell_type": "markdown",
"id": "653eb22f-b09e-4421-8028-05d123fc47a5",
"metadata": {
"id": "653eb22f-b09e-4421-8028-05d123fc47a5"
},
"source": [
"### Data preparation"
]
},
{
"cell_type": "markdown",
"id": "42deb721-f734-43ea-bfef-931733d6379b",
"metadata": {
"id": "42deb721-f734-43ea-bfef-931733d6379b"
},
"source": [
"In this notebook, it is assumed that all audio will be resampled to 16kHz and the data and configuration will be stored under `root_dir` as defined below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "577d4be2-10f8-4050-8f40-b2566ef9dfef",
"metadata": {
"id": "577d4be2-10f8-4050-8f40-b2566ef9dfef"
},
"outputs": [],
"source": [
"# sample rate used throughout the notebook\n",
"sample_rate = 16000\n",
"\n",
"# root directory for data preparation, configurations, etc\n",
"root_dir = Path('./')\n",
"\n",
"# data directory\n",
"data_dir = root_dir / 'data'\n",
"data_dir.mkdir(exist_ok=True)\n",
"\n",
"# scripts directory\n",
"scripts_dir = root_dir / 'scripts'\n",
"scripts_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "markdown",
"id": "9f6dc925-04b1-4118-8092-c3107a734f4f",
"metadata": {
"id": "9f6dc925-04b1-4118-8092-c3107a734f4f"
},
"source": [
"Create dictionary with paths for all of the manifests files which will be stored under `data_dir`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "da232926-a529-4409-b983-8a921c578c91",
"metadata": {
"id": "da232926-a529-4409-b983-8a921c578c91"
},
"outputs": [],
"source": [
"dataset_manifest = {\n",
" 'speech_train': data_dir / 'libri_cuts_train.jsonl.gz',\n",
" 'speech_val': data_dir / 'libri_cuts_val.jsonl.gz',\n",
" 'noise_train': data_dir / 'demand_cuts_train.jsonl.gz',\n",
" 'noise_val': data_dir / 'demand_cuts_val.jsonl.gz',\n",
" 'rir_train': data_dir / 'rir_recordings_train.jsonl.gz',\n",
" 'rir_val': data_dir / 'rir_recordings_val.jsonl.gz',\n",
" 'noisy_val': data_dir / 'noisy_cuts_val.jsonl.gz'\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "d0e19997-b151-47f5-b19b-97d8c7e65e2e",
"metadata": {
"id": "d0e19997-b151-47f5-b19b-97d8c7e65e2e"
},
"source": [
"In this tutorial, a subset of LibriSpeech dataset [2] will be downloaded and used as the speech material.\n",
"\n",
"To use a dataset with the Lhotse dataloader, we need to create manifest files from Lhotse cuts (refer to [3] for the details). In this cell, we first download and prepare the LibriSpeech dataset in a Lhotse format and then save it as manifest files for training and validation sets. Note that the target recording in the speech enhancement task is the original (unchanged) clean speech signal, which is defined under the custom field \"target_recording\" in the cuts."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d4156a68-1906-4a4c-9ecf-cab5d79eb13b",
"metadata": {
"id": "d4156a68-1906-4a4c-9ecf-cab5d79eb13b"
},
"outputs": [],
"source": [
"libri_variant = 'mini_librispeech'\n",
"speech_dir = data_dir / 'speech'\n",
"\n",
"libri_root = download_librispeech(speech_dir, dataset_parts=libri_variant)\n",
"\n",
"# Use script from Lhotse to prepate Librispeech dataset to Lhotse format\n",
"libri = prepare_librispeech(\n",
" libri_root, dataset_parts=libri_variant,\n",
")\n",
"cuts_train = CutSet.from_manifests(**libri[\"train-clean-5\"]).trim_to_supervisions()\n",
"cuts_val = CutSet.from_manifests(**libri[\"dev-clean-2\"]).trim_to_supervisions()\n",
"\n",
"# Save the manifest with a custom \"target_recording\"\n",
"with CutSet.open_writer(dataset_manifest['speech_train']) as writer:\n",
" for cut in cuts_train:\n",
" cut.target_recording = cut.recording\n",
" writer.write(cut)\n",
"\n",
"with CutSet.open_writer(dataset_manifest['speech_val']) as writer:\n",
" for cut in cuts_val:\n",
" cut.target_recording = cut.recording\n",
" writer.write(cut)"
]
},
{
"cell_type": "markdown",
"id": "6f5a0eb7-b913-4323-a8ce-1da77521730b",
"metadata": {
"id": "6f5a0eb7-b913-4323-a8ce-1da77521730b"
},
"source": [
"During the training phase, noise data will be used for online augmentation by mixing it with the downloaded speech on-the-fly. During the validation and test phases, the noise will be used to create fixed sets.\n",
"\n",
"The following cell will download and prepare the noise data using a subset of the DEMAND dataset [4]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ba4dcfe4-5614-438d-9c2b-73c0fbf3bc97",
"metadata": {
"id": "ba4dcfe4-5614-438d-9c2b-73c0fbf3bc97"
},
"outputs": [],
"source": [
"noise_dir = data_dir / 'noise'\n",
"noise_data_set = 'STRAFFIC,PSTATION'\n",
"\n",
"# Copy script\n",
"get_demand_script = os.path.join(scripts_dir, 'get_demand_data.py')\n",
"if not os.path.exists(get_demand_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_demand_data.py\n",
"\n",
"if not noise_dir.is_dir():\n",
" noise_dir.mkdir(exist_ok=True)\n",
" !python {get_demand_script} --data_root={noise_dir} --data_sets={noise_data_set}\n",
"else:\n",
" print('Noise directory already exists in:', noise_dir)\n",
"\n",
"noise_dir = data_dir / 'noise'\n",
"demand_recordings = RecordingSet.from_dir(noise_dir, pattern='*.wav')\n",
"\n",
"demand_cuts = CutSet.from_manifests(recordings=demand_recordings)\n",
"shuffled_demand_cuts = demand_cuts.shuffle()\n",
"\n",
"demand_cuts_train = shuffled_demand_cuts.subset(last=len(shuffled_demand_cuts)-3)\n",
"demand_cuts_val = shuffled_demand_cuts.subset(first=3)\n",
"\n",
"demand_cuts_train.to_file(dataset_manifest['noise_train'])\n",
"demand_cuts_val.to_file(dataset_manifest['noise_val'])"
]
},
{
"cell_type": "markdown",
"id": "f9429dd5-fedc-486d-92f1-5a0255d600e6",
"metadata": {
"id": "f9429dd5-fedc-486d-92f1-5a0255d600e6"
},
"source": [
"The following cell will download and prepare a simulated subset from room impulse responses dataset, described in the following paper [5]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3e83f5d2-26cd-44cb-9efe-4eafcd709bfe",
"metadata": {
"id": "3e83f5d2-26cd-44cb-9efe-4eafcd709bfe"
},
"outputs": [],
"source": [
"rir_recordings = RecordingSet()\n",
"rir_raw_dir = download_rir_noise(data_dir)\n",
"rirs = prepare_rir_noise(rir_raw_dir, parts=[\"sim_rir\"])\n",
"rir_recordings = rirs[\"sim_rir\"][\"recordings\"]\n",
"shuffled_rir_recordings = rir_recordings.shuffle()\n",
"\n",
"rir_val_part = int(len(rir_recordings) * 0.1)\n",
"rir_train_part = len(rir_recordings) - rir_val_part\n",
"\n",
"rir_recordings_train = shuffled_rir_recordings.subset(last=rir_train_part)\n",
"rir_recordings_val = shuffled_rir_recordings.subset(first=rir_val_part)\n",
"\n",
"rir_recordings_train.to_file(dataset_manifest['rir_train'])\n",
"rir_recordings_val.to_file(dataset_manifest['rir_val'])"
]
},
{
"cell_type": "markdown",
"id": "b568a03e-57a1-412f-a3a2-85af3a1e3a20",
"metadata": {
"id": "b568a03e-57a1-412f-a3a2-85af3a1e3a20"
},
"source": [
"For this tutorial, a single-channel noisy validation set is constructed by adding speech and noise.\n",
"\n",
"The following block will use based on Lhotse data loader from NeMo to create fixed noisy validation set and save it do `data/val` folder."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "62dc5755-ccc5-4eb8-80d4-d59b418406a3",
"metadata": {
"id": "62dc5755-ccc5-4eb8-80d4-d59b418406a3"
},
"outputs": [],
"source": [
"# Create the cofing for the Lhotse data loader\n",
"val_noise_config = {\n",
"'cuts_path': dataset_manifest['speech_val'].as_posix(), # path to Lhotse cuts manifest with speech signals for augmentation\n",
"'sample_rate': 16000,\n",
"'batch_size': 1,\n",
"'rir_enabled': True, # enable room impulse response augmentation\n",
"'rir_path': dataset_manifest['rir_val'].as_posix(), # path to Lhotse recordings manifest with room impulse response signals\n",
"'rir_prob': 1.0, # probability of applying RIR augmentation\n",
"'noise_path': dataset_manifest['noise_val'].as_posix(), # path to Lhotse cuts manifest with noise signals\n",
"'noise_mix_prob': 1.0, # probability of applying noise augmentation\n",
"'noise_snr': (0, 20), # range of speech-to-noise ratio for the noise augmentation\n",
"'shuffle': False\n",
"}\n",
"\n",
"# Instantiate the data loader\n",
"dl = get_lhotse_dataloader_from_config(\n",
"OmegaConf.create(val_noise_config), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()\n",
")\n",
"\n",
"# Define number of samples for the validation set\n",
"num_examples = 100\n",
"print(f'Get {num_examples} samples for the validation set')\n",
"samples = [sample for sample in islice(dl, num_examples)]\n",
"\n",
"\n",
"# Create folders for saving noisy (input) and clean (target) samples\n",
"val_dir = data_dir / 'val'\n",
"val_noisy_dir = val_dir / 'noisy'\n",
"val_clean_dir = val_dir / 'clean'\n",
"\n",
"val_dir.mkdir(exist_ok=True)\n",
"val_noisy_dir.mkdir(exist_ok =True)\n",
"val_clean_dir.mkdir(exist_ok=True)\n",
"\n",
"val_noisy_basename = 'val_noisy_fileid'\n",
"val_clean_basename = 'val_clean_fileid'\n",
"\n",
"with CutSet.open_writer(dataset_manifest['noisy_val']) as writer:\n",
" for n, sample in enumerate(samples):\n",
" noisy, clean = sample['input_signal'].numpy()[0], sample['target_signal'].numpy()[0]\n",
" #Save\n",
" sf.write(val_noisy_dir / f'{val_noisy_basename}_{str(n)}.wav', noisy, samplerate=val_noise_config['sample_rate'])\n",
" sf.write(val_clean_dir / f'{val_clean_basename}_{str(n)}.wav', clean, samplerate=val_noise_config['sample_rate'])\n",
" noisy_rec = Recording.from_file(val_noisy_dir / f'{val_noisy_basename}_{str(n)}.wav')\n",
" clean_rec = Recording.from_file(val_clean_dir / f'{val_clean_basename}_{str(n)}.wav')\n",
"\n",
" val_cut = MonoCut(id=noisy_rec.id,\n",
" start=0,\n",
" duration=noisy_rec.duration,\n",
" channel=0,\n",
" recording=noisy_rec)\n",
" val_cut.target_recording = clean_rec\n",
" writer.write(val_cut)"
]
},
{
"cell_type": "markdown",
"id": "f108c5d2-c87b-47a7-8c2d-d363a9234abb",
"metadata": {
"id": "f108c5d2-c87b-47a7-8c2d-d363a9234abb"
},
"source": [
"### Model configuration\n",
"\n",
"Here, a simple encoder-mask-decoder model will be used to process the noisy input signal and produce an enhanced output signal.\n",
"\n",
"In general, an encoder-mask-decoder model can be configured using `EncMaskDecAudioToAudioModel` class, which is depicted in the following block diagram."
]
},
{
"cell_type": "markdown",
"id": "c0ff2bff-1637-4a17-85b2-c92307a4d8d7",
"metadata": {
"id": "c0ff2bff-1637-4a17-85b2-c92307a4d8d7"
},
"source": [
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.18.0/encmaskdecoder_model.png\" alt=\"encmaskdecoder_model\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"id": "470b0a50-fe6c-4ded-9d27-64db20f83cad",
"metadata": {
"id": "470b0a50-fe6c-4ded-9d27-64db20f83cad"
},
"source": [
"The model structure can briefly be described as follows:\n",
"* Input to the model is a time-domain signal.\n",
"* Encoder transforms the input signal to the analysis domain.\n",
"* Mask estimator estimates a mask used to generate the output signal.\n",
"* Mask processor combines the estimated mask and the encoded input to produce the encoded output.\n",
"* Decoder transforms the encoded output into a time-domain signal.\n",
"* Output is a time-domain signal."
]
},
{
"cell_type": "markdown",
"id": "4cf1cfde-913f-42da-a9e0-fdb5ae8e50c5",
"metadata": {
"id": "4cf1cfde-913f-42da-a9e0-fdb5ae8e50c5"
},
"source": [
"For this example, the model will be configured to use a fixed short-time Fourier transform-based encoder and decoder, and the mask will be estimated using a recurrent neural network. The model used here is depicted in the following block diagram."
]
},
{
"cell_type": "markdown",
"id": "8ed37321-2e8b-4d8e-90bf-efc0897517e5",
"metadata": {
"id": "8ed37321-2e8b-4d8e-90bf-efc0897517e5"
},
"source": [
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.18.0/single_output_example_model.png\" alt=\"single_output_example_model\" style=\"width: 1000px;\"/>"
]
},
{
"cell_type": "markdown",
"id": "36596c01-db7d-402f-a5df-0fb9b4641b08",
"metadata": {
"id": "36596c01-db7d-402f-a5df-0fb9b4641b08"
},
"source": [
"In this particular configuration, the model structure can be described as follows:\n",
"* `AudioToSpectrogram` implements the analysis STFT transform.\n",
"* `MaskEstimatorRNN` is a mask estimator using RNNs.\n",
"* `MaskReferenceChannel` is a simple processor which applies the estimated mask on the reference channel. In this tutorial, the input signal has only a single channel, so the reference channel will be set to `0`.\n",
"* `SpectrogramToAudio` implements the synthesis STFT transform."
]
},
{
"cell_type": "markdown",
"id": "03546307-1c4e-4775-a754-276c9a69be5c",
"metadata": {
"id": "03546307-1c4e-4775-a754-276c9a69be5c"
},
"source": [
"The following cell will load and show the default configuration for the model depicted above."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "03952acf-1792-4470-9120-dddf20a9f646",
"metadata": {
"id": "03952acf-1792-4470-9120-dddf20a9f646",
"scrolled": true
},
"outputs": [],
"source": [
"config_dir = root_dir / 'conf'\n",
"config_dir.mkdir(exist_ok=True)\n",
"\n",
"config_path = config_dir / 'masking_with_online_augmentation.yaml'\n",
"\n",
"if not config_path.is_file():\n",
" !wget https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/examples/audio/conf/masking_with_online_augmentation.yaml -P {config_dir.as_posix()}\n",
"\n",
"config = OmegaConf.load(config_path)\n",
"config = OmegaConf.to_container(config, resolve=True)\n",
"config = OmegaConf.create(config)\n",
"\n",
"print('Loaded config')\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"id": "1c872345-3bfd-4933-8c51-dbb62fce790a",
"metadata": {
"id": "1c872345-3bfd-4933-8c51-dbb62fce790a"
},
"source": [
"Training dataset is configured with the following parameters\n",
"* `cuts_path` points to a Lhotse manifest file, containing speech samples\n",
"* `noise_path` poins to a Lhotse manifest file, containing noise samples\n",
"* `noise_mix_prob` defines the probabilty with which noise will be added during training\n",
"* `noise_snr` defines an SNR range for mixing noise samples\n",
"* `rir_enabled` enables room impulse response agmentation\n",
"* `rir_path` points to a Lhotse manifest file, containing RIR samples\n",
"* `rir_prob` defines the probabilty with which RIR will be added during training\n",
" \n",
"For the validation and test sets only `cuts_path` parameter is used since the `val` manifest already contains noisy and clean samples"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e19c4209-b10b-4924-a3de-e792de57c313",
"metadata": {
"id": "e19c4209-b10b-4924-a3de-e792de57c313"
},
"outputs": [],
"source": [
"# Setup training dataset\n",
"config.model.train_ds.cuts_path = dataset_manifest['speech_train'].as_posix() # path to Lhotse cuts manifest with speech signals for augmentation\n",
"config.model.train_ds.noise_path = dataset_manifest['noise_train'].as_posix() # path to Lhotse cuts manifest with noise signals\n",
"config.model.train_ds.noise_mix_prob = 1.0 # probability of applying noise augmentation\n",
"config.model.train_ds.noise_snr = (0, 20) # range of speech-to-noise ratio for the noise augmentation\n",
"config.model.train_ds.rir_enabled = True # enable room impulse response augmentation\n",
"config.model.train_ds.rir_path = dataset_manifest['rir_val'].as_posix() # path to Lhotse recordings manifest with room impulse response signals\n",
"config.model.train_ds.rir_prob = 1.0 # probability of applying RIR augmentation\n",
"\n",
"config.model.validation_ds.cuts_path = dataset_manifest['noisy_val'].as_posix() # fixed noisy validation set\n",
"\n",
"config.model.test_ds.cuts_path = dataset_manifest['noisy_val'].as_posix() # fixed noisy test set\n",
"\n",
"\n",
"print(\"Train dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.train_ds))"
]
},
{
"cell_type": "markdown",
"id": "534aef42-d3a7-4477-8f89-8fe780f2b5f7",
"metadata": {
"id": "534aef42-d3a7-4477-8f89-8fe780f2b5f7"
},
"source": [
"Metrics for validation and test set are configured in the following cell.\n",
"\n",
"In this tutorial, signal-to-distortion ratio (SDR) and scale-invariant SDR from torch metrics are used [5]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "522f1065-fe23-4101-b6a9-0706cc8db389",
"metadata": {
"id": "522f1065-fe23-4101-b6a9-0706cc8db389"
},
"outputs": [],
"source": [
"# Setup metrics to compute on validation and test sets\n",
"metrics = OmegaConf.create({\n",
" 'sisdr': {\n",
" '_target_': 'torchmetrics.audio.ScaleInvariantSignalDistortionRatio',\n",
" },\n",
" 'sdr': {\n",
" '_target_': 'torchmetrics.audio.SignalDistortionRatio',\n",
" }\n",
"})\n",
"config.model.metrics.val = metrics\n",
"config.model.metrics.test = metrics\n",
"\n",
"print(\"Metrics config:\")\n",
"print(OmegaConf.to_yaml(config.model.metrics))"
]
},
{
"cell_type": "markdown",
"id": "4771d02f-3b01-481c-8894-bb43644a941a",
"metadata": {
"id": "4771d02f-3b01-481c-8894-bb43644a941a"
},
"source": [
"### Trainer configuration\n",
"NeMo models are primarily PyTorch Lightning modules and therefore are entirely compatible with the PyTorch Lightning ecosystem."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2208b407-07e2-4bb0-bfad-2ef488acf217",
"metadata": {
"id": "2208b407-07e2-4bb0-bfad-2ef488acf217"
},
"outputs": [],
"source": [
"print(\"Trainer config:\")\n",
"print(OmegaConf.to_yaml(config.trainer))"
]
},
{
"cell_type": "markdown",
"id": "f2dac35a-b487-4037-907d-326f321564ca",
"metadata": {
"id": "f2dac35a-b487-4037-907d-326f321564ca"
},
"source": [
"We can modify some trainer configs for this tutorial.\n",
"Most importantly, the number of epochs is set to a small value, to limit the runtime for the purpose of this example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2c18dbd2-ad22-4eae-912a-56d13aa74a6b",
"metadata": {
"id": "2c18dbd2-ad22-4eae-912a-56d13aa74a6b"
},
"outputs": [],
"source": [
"# Checks if we have GPU available and uses it\n",
"accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
"config.trainer.devices = 1\n",
"config.trainer.accelerator = accelerator\n",
"\n",
"# Reduces maximum number of epochs for quick demonstration\n",
"config.trainer.max_epochs = 30\n",
"\n",
"# Remove distributed training flags\n",
"config.trainer.strategy = 'auto'\n",
"\n",
"# Instantiate the trainer\n",
"trainer = pl.Trainer(**config.trainer)"
]
},
{
"cell_type": "markdown",
"id": "c1198508-7fb5-4ba7-a4e5-ad15153037a1",
"metadata": {
"id": "c1198508-7fb5-4ba7-a4e5-ad15153037a1"
},
"source": [
"### Experiment manager\n",
"\n",
"NeMo has an experiment manager that handles logging and checkpointing."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "72642c8c-97f5-4b6c-9e63-030c7a07810c",
"metadata": {
"id": "72642c8c-97f5-4b6c-9e63-030c7a07810c",
"scrolled": true
},
"outputs": [],
"source": [
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"# The exp_dir provides a path to the current experiment for easy access\n",
"\n",
"print(\"Experiment directory:\")\n",
"print(exp_dir)"
]
},
{
"cell_type": "markdown",
"id": "00562a21-0c0f-4c53-b89e-344431e75a42",
"metadata": {
"id": "00562a21-0c0f-4c53-b89e-344431e75a42"
},
"source": [
"### Model instantiation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b73c3de-6ded-491d-8026-8a093f7ffc16",
"metadata": {
"id": "0b73c3de-6ded-491d-8026-8a093f7ffc16",
"scrolled": true
},
"outputs": [],
"source": [
"from nemo.collections.audio.models.enhancement import EncMaskDecAudioToAudioModel\n",
"\n",
"enhancement_model = EncMaskDecAudioToAudioModel(cfg=config.model, trainer=trainer)"
]
},
{
"cell_type": "markdown",
"id": "a8f4373e-6aca-41d3-a69d-5bef53a88f20",
"metadata": {
"id": "a8f4373e-6aca-41d3-a69d-5bef53a88f20"
},
"source": [
"### Training\n",
"Create a Tensorboard visualization to monitor progress"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a253ce7c-6cca-441f-818e-1c8a2f00a9f5",
"metadata": {
"id": "a253ce7c-6cca-441f-818e-1c8a2f00a9f5"
},
"outputs": [],
"source": [
"try:\n",
" from google import colab\n",
" COLAB_ENV = True\n",
"except (ImportError, ModuleNotFoundError):\n",
" COLAB_ENV = False\n",
"\n",
"# Load the TensorBoard notebook extension\n",
"if COLAB_ENV:\n",
" %load_ext tensorboard\n",
" %tensorboard --logdir {exp_dir}\n",
"else:\n",
" print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
]
},
{
"cell_type": "markdown",
"id": "9183a35d-8332-49f6-96b3-3d14dd70ff4e",
"metadata": {
"id": "9183a35d-8332-49f6-96b3-3d14dd70ff4e"
},
"source": [
"Training can be started using `trainer.fit`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0eab4e7e-3004-49fa-a09b-10dfd12ba294",
"metadata": {
"id": "0eab4e7e-3004-49fa-a09b-10dfd12ba294",
"scrolled": true
},
"outputs": [],
"source": [
"trainer.fit(enhancement_model)"
]
},
{
"cell_type": "markdown",
"id": "af906f16-626c-4ca4-a465-fb0cca8d514a",
"metadata": {
"id": "af906f16-626c-4ca4-a465-fb0cca8d514a"
},
"source": [
"After the training is completed, the configured metrics can be easily computed on the test set as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "589af425-18b2-4445-8c88-d3161e9928c5",
"metadata": {
"id": "589af425-18b2-4445-8c88-d3161e9928c5"
},
"outputs": [],
"source": [
"trainer.test(enhancement_model, ckpt_path=None)"
]
},
{
"cell_type": "markdown",
"id": "cce96446-a969-4c6f-a865-968ee61119a8",
"metadata": {
"id": "cce96446-a969-4c6f-a865-968ee61119a8"
},
"source": [
"### Inference\n",
"\n",
"The following cell provides an example of inference on an single audio file.\n",
"For simplicity, the audio file information is taken from the test dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9f29ade5-424e-4a7a-81b7-dc3c5cc90bf7",
"metadata": {
"id": "9f29ade5-424e-4a7a-81b7-dc3c5cc90bf7"
},
"outputs": [],
"source": [
"# Load 10 samples from test_dataloader\n",
"samples = [sample for sample in islice(enhancement_model.test_dataloader(), 10)]\n",
"\n",
"# Different sample can be used via list index\n",
"sample = samples[0]\n",
"\n",
"noisy_tensor = sample['input_signal']\n",
"speech_tensor = sample['target_signal']\n",
"\n",
"# Get the one-dimentional numpy signals for the plotting audio files and metrics calculation\n",
"noisy_signal = noisy_tensor.squeeze(0).numpy()\n",
"speech_signal = speech_tensor.squeeze(0).numpy()\n",
"\n",
"\n",
"# Move to device\n",
"device = 'cuda' if accelerator == 'gpu' else 'cpu'\n",
"enhancement_model = enhancement_model.to(device)\n",
"\n",
"# Process using the model\n",
"with torch.no_grad():\n",
" output_tensor, _ = enhancement_model(input_signal=noisy_tensor.unsqueeze(1).cuda())\n",
"output_signal = output_tensor[0][0].detach().cpu().numpy()"
]
},
{
"cell_type": "markdown",
"id": "5110ee6c-3939-4d70-a4d6-6422de31da51",
"metadata": {
"id": "5110ee6c-3939-4d70-a4d6-6422de31da51"
},
"source": [
"Signals can be easily plotted and signal metrics can be calculated for the given example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0275a032-969c-4e31-b8b5-8a61cb1f82ee",
"metadata": {
"id": "0275a032-969c-4e31-b8b5-8a61cb1f82ee"
},
"outputs": [],
"source": [
"# Show noisy and clean signals\n",
"show_metrics(signal=noisy_signal, reference=speech_signal, tag='Noisy signal', sample_rate=sample_rate)\n",
"show_metrics(signal=output_signal, reference=speech_signal, tag='Output signal', sample_rate=sample_rate)\n",
"\n",
"# Show signals\n",
"show_signal(speech_signal, tag='Speech signal')\n",
"show_signal(noisy_signal, tag='Noisy signal')\n",
"show_signal(output_signal, tag='Output signal')\n",
"\n",
"# Play audio\n",
"print('Speech signal')\n",
"ipd.display(ipd.Audio(speech_signal, rate=sample_rate))\n",
"\n",
"print('Noisy signal')\n",
"ipd.display(ipd.Audio(noisy_signal, rate=sample_rate))\n",
"\n",
"print('Output signal')\n",
"ipd.display(ipd.Audio(output_signal, rate=sample_rate))"
]
},
{
"cell_type": "markdown",
"id": "89eb3478-bf82-4fbd-aa7d-0a55edf57f3a",
"metadata": {
"id": "89eb3478-bf82-4fbd-aa7d-0a55edf57f3a"
},
"source": [
"## Next steps\n",
"This is a simple tutorial which can serve as a starting point for prototyping and experimentation with audio-to-audio models.\n",
"A processed audio output can be used, for example, for ASR or TTS.\n",
"\n",
"For more details about NeMo models and applications in in ASR and TTS, we recommend you checkout other tutorials next:\n",
"\n",
"* [NeMo fundamentals](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb)\n",
"* [NeMo models](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb)\n",
"* [Speech Recognition](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)\n",
"* [Speech Synthesis](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/tts/Inference_ModelSelect.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "31a162ae-6b9e-40e4-a468-fbbad3360b11",
"metadata": {
"id": "31a162ae-6b9e-40e4-a468-fbbad3360b11"
},
"source": [
"## References\n",
"\n",
"[1] Żelasko, P., Povey, D., Trmal, J., & Khudanpur, S. (2021). Lhotse: a speech data representation library for the modern deep learning ecosystem. https://arxiv.org/abs/2110.12561\n",
"\n",
"[2] V. Panayotov, G. Chen, D. Povery, S. Khudanpur, \"LibriSpeech: An ASR corpus based on public domain audio books,\" ICASSP 2015\n",
"\n",
"[3] Lhotse documentation, https://lhotse.readthedocs.io/\n",
"\n",
"[4] J. Thieman, N. Ito, V. Emmanuel, \"DEMAND: collection of multi-channel recordings of acoustic noise in diverse environments,\" ICA 2013\n",
"\n",
"[5] T. Ko, V. Peddinti, D. Povey, M. L. Seltzer and S. Khudanpur, \"A study on data augmentation of reverberant speech for robust speech recognition,\" 2017 ICASSP\n",
"\n",
"[6] https://github.com/Lightning-AI/torchmetrics"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}