chore: import upstream snapshot with attribution
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,318 @@
|
||||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Script to compute metrics for a given audio-to-audio model for a given manifest file for some dataset.
|
||||
The manifest file must include path to input audio and path to target (ground truth) audio.
|
||||
|
||||
Note: This scripts depends on the `process_audio.py` script, and therefore both scripts should be
|
||||
located in the same directory during execution.
|
||||
|
||||
# Arguments
|
||||
|
||||
<< All arguments of `process_audio.py` are inherited by this script, so please refer to `process_audio.py`
|
||||
for full list of arguments >>
|
||||
|
||||
dataset_manifest: Required - path to dataset JSON manifest file (in NeMo format)
|
||||
output_dir: Optional - output directory where the processed audio will be saved
|
||||
metrics: Optional - list of metrics to evaluate. Defaults to [sdr,estoi]
|
||||
sample_rate: Optional - sample rate for loaded audio. Defaults to 16kHz.
|
||||
only_score_manifest: Optional - If set, processing will be skipped and it is assumed the processed audio is available in dataset_manifest
|
||||
|
||||
# Usage
|
||||
|
||||
## To score a dataset with a manifest file that contains the input audio which needs to be processed and target audio
|
||||
|
||||
python audio_to_audio_eval.py \
|
||||
model_path=null \
|
||||
pretrained_model=null \
|
||||
dataset_manifest=<Mandatory: path to a dataset manifest file> \
|
||||
output_dir=<Optional: Directory where processed audio will be saved> \
|
||||
processed_channel_selector=<Optional: list of channels to select from the processed audio file> \
|
||||
target_key=<Optional: key for the target audio in the dataset manifest. Default: target_audio_filepath> \
|
||||
target_channel_selector=<Optional: list of channels to select from the target audio file> \
|
||||
metrics=<Optional: list of metrics to evaluate. Defaults to [sdr,estoi]>
|
||||
batch_size=32 \
|
||||
amp=True
|
||||
|
||||
## To score a manifest file which has been previously processed and contains both processed audio and target audio
|
||||
|
||||
python audio_to_audio_eval.py \
|
||||
dataset_manifest=<Mandatory: path to a dataset manifest file> \
|
||||
processed_key=<Optional: key for the target audio in the dataset manifest. Default: processed_audio_filepath>
|
||||
processed_channel_selector=<Optional: list of channels to select from the processed audio file> \
|
||||
target_key=<Optional: key for the target audio in the dataset manifest. Default: target_audio_filepath> \
|
||||
target_channel_selector=<Optional: list of channels to select from the target audio file> \
|
||||
metrics=<Optional: list of metrics to evaluate. Defaults to [sdr,estoi]>
|
||||
batch_size=32 \
|
||||
amp=True
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field, is_dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
import process_audio
|
||||
import torch
|
||||
from omegaconf import OmegaConf, open_dict
|
||||
from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
|
||||
from torchmetrics.audio.sdr import ScaleInvariantSignalDistortionRatio, SignalDistortionRatio
|
||||
from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility
|
||||
from tqdm import tqdm
|
||||
|
||||
from nemo.collections.audio.data import audio_to_audio_dataset
|
||||
from nemo.collections.audio.data.audio_to_audio_lhotse import LhotseAudioToTargetDataset
|
||||
from nemo.collections.audio.metrics import AudioMetricWrapper, SquimMOSMetric, SquimObjectiveMetric
|
||||
from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
|
||||
from nemo.collections.common.parts.preprocessing import manifest
|
||||
from nemo.core.config import hydra_runner
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioEvaluationConfig(process_audio.ProcessConfig):
|
||||
# Processed audio config
|
||||
processed_channel_selector: Optional[List] = None
|
||||
processed_key: str = 'processed_audio_filepath'
|
||||
|
||||
# Target audio configs
|
||||
target_dataset_dir: Optional[str] = None # If not provided, defaults to dirname(cfg.dataset_manifest)
|
||||
target_channel_selector: Optional[List] = None
|
||||
target_key: str = 'target_audio_filepath'
|
||||
|
||||
# Sample rate for audio evaluation
|
||||
sample_rate: int = 16000
|
||||
|
||||
# Score an existing manifest without running processing
|
||||
only_score_manifest: bool = False
|
||||
|
||||
# Metrics to calculate
|
||||
metrics: List[str] = field(default_factory=lambda: ['sdr', 'estoi'])
|
||||
|
||||
# Return metric values for each example
|
||||
return_values_per_example: bool = False
|
||||
|
||||
|
||||
def get_evaluation_dataloader(config):
|
||||
"""Prepare a dataloader for evaluation."""
|
||||
if config.get("use_lhotse", False):
|
||||
return get_lhotse_dataloader_from_config(
|
||||
config, global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
||||
)
|
||||
|
||||
dataset = audio_to_audio_dataset.get_audio_to_target_dataset(config=config)
|
||||
|
||||
return torch.utils.data.DataLoader(
|
||||
dataset=dataset,
|
||||
batch_size=config['batch_size'],
|
||||
collate_fn=dataset.collate_fn,
|
||||
drop_last=config.get('drop_last', False),
|
||||
shuffle=False,
|
||||
num_workers=config.get('num_workers', min(config['batch_size'], os.cpu_count() - 1)),
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
|
||||
def get_metrics(cfg: AudioEvaluationConfig):
|
||||
"""Prepare a dictionary with metrics."""
|
||||
available_metrics = [
|
||||
'sdr',
|
||||
'sisdr',
|
||||
'stoi',
|
||||
'estoi',
|
||||
'pesq',
|
||||
'squim_mos',
|
||||
'squim_stoi',
|
||||
'squim_pesq',
|
||||
'squim_si_sdr',
|
||||
]
|
||||
|
||||
metrics = dict()
|
||||
for name in sorted(set(cfg.metrics)):
|
||||
name = name.lower()
|
||||
if name == 'sdr':
|
||||
metric = AudioMetricWrapper(metric=SignalDistortionRatio())
|
||||
elif name == 'sisdr':
|
||||
metric = AudioMetricWrapper(metric=ScaleInvariantSignalDistortionRatio())
|
||||
elif name == 'stoi':
|
||||
metric = AudioMetricWrapper(metric=ShortTimeObjectiveIntelligibility(fs=cfg.sample_rate, extended=False))
|
||||
elif name == 'estoi':
|
||||
metric = AudioMetricWrapper(metric=ShortTimeObjectiveIntelligibility(fs=cfg.sample_rate, extended=True))
|
||||
elif name == 'pesq':
|
||||
metric = AudioMetricWrapper(metric=PerceptualEvaluationSpeechQuality(fs=cfg.sample_rate, mode='wb'))
|
||||
elif name == 'squim_mos':
|
||||
metric = AudioMetricWrapper(metric=SquimMOSMetric(fs=cfg.sample_rate))
|
||||
elif name == 'squim_stoi':
|
||||
metric = AudioMetricWrapper(metric=SquimObjectiveMetric(metric='stoi', fs=cfg.sample_rate))
|
||||
elif name == 'squim_pesq':
|
||||
metric = AudioMetricWrapper(metric=SquimObjectiveMetric(metric='pesq', fs=cfg.sample_rate))
|
||||
elif name == 'squim_si_sdr':
|
||||
metric = AudioMetricWrapper(metric=SquimObjectiveMetric(metric='si_sdr', fs=cfg.sample_rate))
|
||||
else:
|
||||
raise ValueError(f'Unexpected metric: {name}. Currently available metrics: {available_metrics}')
|
||||
|
||||
metrics[name] = metric
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
@hydra_runner(config_name="AudioEvaluationConfig", schema=AudioEvaluationConfig)
|
||||
def main(cfg: AudioEvaluationConfig):
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
if is_dataclass(cfg):
|
||||
cfg = OmegaConf.structured(cfg)
|
||||
|
||||
if cfg.audio_dir is not None:
|
||||
raise RuntimeError(
|
||||
"Evaluation script requires ground truth audio to be passed via a manifest file. "
|
||||
"If manifest file is available, submit it via `dataset_manifest` argument."
|
||||
)
|
||||
|
||||
if not os.path.exists(cfg.dataset_manifest):
|
||||
raise FileNotFoundError(f'The dataset manifest file could not be found at path : {cfg.dataset_manifest}')
|
||||
|
||||
if cfg.target_dataset_dir is None:
|
||||
# Assume the target data is available in the same directory as the input data
|
||||
cfg.target_dataset_dir = os.path.dirname(cfg.dataset_manifest)
|
||||
elif not os.path.isdir(cfg.target_dataset_dir):
|
||||
raise FileNotFoundError(f'Target dataset dir could not be found at path : {cfg.target_dataset_dir}')
|
||||
|
||||
# Setup metrics
|
||||
metrics = get_metrics(cfg)
|
||||
|
||||
if cfg.return_values_per_example and cfg.batch_size > 1:
|
||||
raise ValueError('return_example_values is only supported for batch_size=1.')
|
||||
|
||||
# Processing
|
||||
if not cfg.only_score_manifest:
|
||||
# Process audio using the configured model and save in the output directory
|
||||
process_cfg = process_audio.main(cfg) # type: ProcessConfig
|
||||
|
||||
# Release GPU memory if it was used during transcription
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
logging.info('Finished processing audio.')
|
||||
else:
|
||||
# Score the input manifest, no need to run a model
|
||||
cfg.output_filename = cfg.dataset_manifest
|
||||
process_cfg = cfg
|
||||
|
||||
# Evaluation
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
# Prepare a temporary manifest with processed audio and target
|
||||
temporary_manifest_filepath = os.path.join(tmp_dir, 'manifest.json')
|
||||
|
||||
num_files = 0
|
||||
|
||||
with (
|
||||
open(process_cfg.output_filename, 'r') as f_processed,
|
||||
open(temporary_manifest_filepath, 'w', encoding='utf-8') as f_tmp,
|
||||
):
|
||||
for line_processed in f_processed:
|
||||
data_processed = json.loads(line_processed)
|
||||
|
||||
if cfg.processed_key not in data_processed:
|
||||
raise ValueError(
|
||||
f'Processed key {cfg.processed_key} not found in manifest: {process_cfg.output_filename}.'
|
||||
)
|
||||
|
||||
if cfg.target_key not in data_processed:
|
||||
raise ValueError(
|
||||
f'Target key {cfg.target_key} not found in manifest: {process_cfg.output_filename}.'
|
||||
)
|
||||
|
||||
item = {
|
||||
'processed': manifest.get_full_path(
|
||||
audio_file=data_processed[cfg.processed_key], manifest_file=process_cfg.output_filename
|
||||
),
|
||||
'target': manifest.get_full_path(
|
||||
audio_file=data_processed[cfg.target_key], data_dir=cfg.target_dataset_dir
|
||||
),
|
||||
'duration': data_processed.get('duration'),
|
||||
}
|
||||
|
||||
# Double-check files exist
|
||||
for key in ['processed', 'target']:
|
||||
if not os.path.isfile(item[key]):
|
||||
raise ValueError(f'File for key "{key}" not found at: {item[key]}.\nCurrent item: {item}')
|
||||
|
||||
# Warn if we're comparing the same files
|
||||
if item['target'] == item['processed']:
|
||||
logging.warning('Using the same file as processed and target: %s', item['target'])
|
||||
|
||||
# Write the entry in the temporary manifest file
|
||||
f_tmp.write(json.dumps(item) + '\n')
|
||||
|
||||
num_files += 1
|
||||
|
||||
if cfg.max_utts is not None and num_files >= cfg.max_utts:
|
||||
logging.info('Reached max_utts: %s', cfg.max_utts)
|
||||
break
|
||||
|
||||
# Prepare dataloader
|
||||
config = {
|
||||
'manifest_filepath': temporary_manifest_filepath,
|
||||
'sample_rate': cfg.sample_rate,
|
||||
'input_key': 'processed',
|
||||
'input_channel_selector': cfg.processed_channel_selector,
|
||||
'target_key': 'target',
|
||||
'target_channel_selector': cfg.target_channel_selector,
|
||||
'batch_size': min(cfg.batch_size, num_files),
|
||||
'num_workers': cfg.num_workers,
|
||||
}
|
||||
temporary_dataloader = get_evaluation_dataloader(config)
|
||||
|
||||
metrics_value_per_example = defaultdict(list)
|
||||
|
||||
# Calculate metrics
|
||||
for eval_batch in tqdm(temporary_dataloader, desc='Evaluating'):
|
||||
processed_signal, processed_length, target_signal, target_length = eval_batch
|
||||
|
||||
if not torch.equal(processed_length, target_length):
|
||||
raise RuntimeError(f'Length mismatch.')
|
||||
|
||||
for name, metric in metrics.items():
|
||||
value = metric(preds=processed_signal, target=target_signal, input_length=target_length)
|
||||
if cfg.return_values_per_example:
|
||||
metrics_value_per_example[name].append(value.item())
|
||||
|
||||
# Convert to a dictionary with name: value
|
||||
metrics_value = {name: metric.compute().item() for name, metric in metrics.items()}
|
||||
|
||||
logging.info('Finished running evaluation.')
|
||||
|
||||
# Show results
|
||||
logging.info('Summary\n')
|
||||
logging.info('Data')
|
||||
logging.info('\tmanifest: %s', cfg.output_filename)
|
||||
logging.info('\ttarget_dataset_dir: %s', cfg.target_dataset_dir)
|
||||
logging.info('\tnum_files: %s', num_files)
|
||||
logging.info('Metrics')
|
||||
for name, value in metrics_value.items():
|
||||
logging.info('\t%10s: \t%6.2f', name, value)
|
||||
|
||||
# Inject the metric name and score into the config, and return the entire config
|
||||
with open_dict(cfg):
|
||||
cfg.metrics_value = metrics_value
|
||||
cfg.metrics_value_per_example = dict(metrics_value_per_example)
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
@@ -0,0 +1,116 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
# Training the model
|
||||
|
||||
Basic run (on CPU for 50 epochs):
|
||||
python examples/audio/audio_to_audio_train.py \
|
||||
# (Optional: --config-path=<path to dir of configs> --config-name=<name of config without .yaml>) \
|
||||
model.train_ds.manifest_filepath="<path to manifest file>" \
|
||||
model.validation_ds.manifest_filepath="<path to manifest file>" \
|
||||
trainer.devices=1 \
|
||||
trainer.accelerator='cpu' \
|
||||
trainer.max_epochs=50
|
||||
|
||||
PyTorch Lightning Trainer arguments and args of the model and the optimizer can be added or overriden from CLI
|
||||
"""
|
||||
from enum import Enum
|
||||
|
||||
import lightning.pytorch as pl
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.audio.models.enhancement import (
|
||||
EncMaskDecAudioToAudioModel,
|
||||
FlowMatchingAudioToAudioModel,
|
||||
PredictiveAudioToAudioModel,
|
||||
SchroedingerBridgeAudioToAudioModel,
|
||||
ScoreBasedGenerativeAudioToAudioModel,
|
||||
)
|
||||
from nemo.collections.audio.models.maxine import BNR2
|
||||
from nemo.core.config import hydra_runner
|
||||
from nemo.utils import logging
|
||||
from nemo.utils.exp_manager import exp_manager
|
||||
|
||||
|
||||
class ModelType(str, Enum):
|
||||
"""Enumeration with the available model types."""
|
||||
|
||||
MaskBased = 'mask_based'
|
||||
Predictive = 'predictive'
|
||||
ScoreBased = 'score_based'
|
||||
SchroedingerBridge = 'schroedinger_bridge'
|
||||
FlowMatching = 'flow_matching'
|
||||
BNR2 = 'bnr'
|
||||
|
||||
|
||||
def get_model_class(model_type: ModelType):
|
||||
"""Get model class for a given model type."""
|
||||
if model_type == ModelType.MaskBased:
|
||||
return EncMaskDecAudioToAudioModel
|
||||
elif model_type == ModelType.Predictive:
|
||||
return PredictiveAudioToAudioModel
|
||||
elif model_type == ModelType.ScoreBased:
|
||||
return ScoreBasedGenerativeAudioToAudioModel
|
||||
elif model_type == ModelType.SchroedingerBridge:
|
||||
return SchroedingerBridgeAudioToAudioModel
|
||||
elif model_type == ModelType.FlowMatching:
|
||||
return FlowMatchingAudioToAudioModel
|
||||
elif model_type == ModelType.BNR2:
|
||||
return BNR2
|
||||
else:
|
||||
raise ValueError(f'Unknown model type: {model_type}')
|
||||
|
||||
|
||||
@hydra_runner(config_path="./conf", config_name="masking")
|
||||
def main(cfg):
|
||||
logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg, resolve=True)}')
|
||||
|
||||
trainer = pl.Trainer(**cfg.trainer)
|
||||
exp_manager(trainer, cfg.get("exp_manager", None))
|
||||
|
||||
# Get model class
|
||||
model_type = cfg.model.get('type')
|
||||
if model_type is None:
|
||||
model_type = ModelType.MaskBased
|
||||
logging.warning('model_type not found in config. Using default: %s', model_type)
|
||||
|
||||
logging.info('Get class for model type: %s', model_type)
|
||||
model_class = get_model_class(model_type)
|
||||
|
||||
logging.info('Instantiate model %s', model_class.__name__)
|
||||
model = model_class(cfg=cfg.model, trainer=trainer)
|
||||
|
||||
logging.info('Initialize the weights of the model from another model, if provided via config')
|
||||
model.maybe_init_from_pretrained_checkpoint(cfg)
|
||||
|
||||
# Train the model
|
||||
trainer.fit(model)
|
||||
|
||||
# Run on test data, if available
|
||||
if hasattr(cfg.model, 'test_ds'):
|
||||
if trainer.is_global_zero:
|
||||
# Destroy the current process group and let the trainer initialize it again with a single device.
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
# Run test on a single device
|
||||
trainer = pl.Trainer(devices=1, accelerator=cfg.trainer.accelerator)
|
||||
if model.prepare_test(trainer):
|
||||
trainer.test(model)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
@@ -0,0 +1,125 @@
|
||||
# This configuration contains the exemplary values for training a multichannel speech enhancement model with a mask-based beamformer.
|
||||
#
|
||||
name: "beamforming"
|
||||
|
||||
model:
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
audio_duration: 4.0 # in seconds, audio segment duration for training
|
||||
random_offset: true # if the file is longer than audio_duration, use random offset to select a subsegment
|
||||
min_duration: ${model.train_ds.audio_duration}
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
batch_size: 1 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
test_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
batch_size: 1 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
mask_estimator:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskEstimatorRNN
|
||||
num_outputs: ${model.num_outputs}
|
||||
num_subbands: 257 # Number of subbands of the input spectrogram
|
||||
num_features: 256 # Number of features at RNN input
|
||||
num_layers: 5 # Number of RNN layers
|
||||
bidirectional: true # Use bi-directional RNN
|
||||
|
||||
mask_processor:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskBasedBeamformer # Mask-based multi-channel processing
|
||||
ref_channel: 0 # Reference channel for the output
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.SDRLoss
|
||||
scale_invariant: true # Use scale-invariant SDR
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sdr: # output SDR
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
test:
|
||||
sdr_ch0: # SDR on output channel 0
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
channel: 0
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: False # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
create_tensorboard_logger: true
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: "val_loss"
|
||||
mode: "min"
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,146 @@
|
||||
# This configuration contains the exemplary values for training a multichannel speech enhancement model with a mask-based beamformer.
|
||||
#
|
||||
name: beamforming_flex_channels
|
||||
|
||||
model:
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
input_channel_selector: null # load all channels from the input file
|
||||
target_key: target_anechoic_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # load only the first channel from the target file
|
||||
audio_duration: 4.0 # in seconds, audio segment duration for training
|
||||
random_offset: true # if the file is longer than audio_duration, use random offset to select a subsegment
|
||||
min_duration: ${model.train_ds.audio_duration}
|
||||
batch_size: 16 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 16
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
input_channel_selector: null # load all channels from the input file
|
||||
target_key: target_anechoic_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # load only the first channel from the target file
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
channel_augment:
|
||||
_target_: nemo.collections.asr.parts.submodules.multichannel_modules.ChannelAugment
|
||||
num_channels_min: 2 # minimal number of channels selected for each batch
|
||||
num_channels_max: null # max number of channels is determined by the batch size
|
||||
permute_channels: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
|
||||
mask_estimator:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskEstimatorFlexChannels
|
||||
num_outputs: ${model.num_outputs} # number of output masks
|
||||
num_subbands: 257 # number of subbands for the input spectrogram
|
||||
num_blocks: 5 # number of blocks in the model
|
||||
channel_reduction_position: 3 # 0-indexed, apply channel reduction before this block
|
||||
channel_reduction_type: average # channel-wise reduction
|
||||
channel_block_type: transform_average_concatenate # channel block
|
||||
temporal_block_type: conformer_encoder # temporal block
|
||||
temporal_block_num_layers: 5 # number of layers for the temporal block
|
||||
temporal_block_num_heads: 4 # number of heads for the temporal block
|
||||
temporal_block_dimension: 128 # the hidden size of the temporal block
|
||||
mag_reduction: null # channel-wise reduction of magnitude
|
||||
mag_normalization: mean_var # normalization using mean and variance
|
||||
use_ipd: true # use inter-channel phase difference
|
||||
ipd_normalization: mean # mean normalization
|
||||
|
||||
mask_processor:
|
||||
# Mask-based multi-channel processor
|
||||
_target_: nemo.collections.audio.modules.masking.MaskBasedBeamformer
|
||||
filter_type: pmwf # parametric multichannel wiener filter
|
||||
filter_beta: 0.0 # mvdr
|
||||
filter_rank: one
|
||||
ref_channel: max_snr # select reference channel by maximizing estimated SNR
|
||||
ref_hard: 1 # a one-hot reference. If false, a soft estimate across channels is used.
|
||||
ref_hard_use_grad: false # use straight-through gradient when using hard reference
|
||||
ref_subband_weighting: false # use subband weighting for reference estimation
|
||||
num_subbands: ${model.mask_estimator.num_subbands}
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.SDRLoss
|
||||
convolution_invariant: true # convolution-invariant loss
|
||||
sdr_max: 30 # soft threshold for SDR
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sdr_0:
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
channel: 0 # evaluate only on channel 0, if there are multiple outputs
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: 10000
|
||||
warmup_ratio: null
|
||||
min_lr: 1e-6
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: False # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
create_tensorboard_logger: true
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: "val_loss"
|
||||
mode: "min"
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.pyth
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,166 @@
|
||||
name: flow_matching_generative
|
||||
|
||||
model:
|
||||
type: flow_matching
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
p_cond: 0.9 # Proability of feeding the conditional input into the model.
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
max_utts_evaluation_metrics: 500
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
audio_duration: 6.14 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 768
|
||||
random_offset: true
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
log_config:
|
||||
log_tensorboard: true
|
||||
log_wandb: false
|
||||
max_utts: 8
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.transformerunet.SpectrogramTransformerUNet
|
||||
in_channels: 2 # concatenation of single-channel perturbed and noisy
|
||||
out_channels: 1 # single-channel score estimate
|
||||
depth: 24
|
||||
ff_dropout: 0.1
|
||||
time_hidden_dim: 1024
|
||||
|
||||
flow:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.OptimalTransportFlow
|
||||
sigma_start: 1.0
|
||||
sigma_end: 1e-4
|
||||
|
||||
sampler:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.ConditionalFlowMatchingEulerSampler
|
||||
num_steps: 20
|
||||
time_min: 1e-8
|
||||
time_max: 1.0
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss
|
||||
ndim: 4 # loss is calculated on the score in the encoded domain (batch, channel, dimension, time)
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: 5000
|
||||
warmup_ratio: null
|
||||
min_lr: 0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: 0.2
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_pesq
|
||||
mode: max
|
||||
save_top_k: 3
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: test
|
||||
project: gense
|
||||
@@ -0,0 +1,169 @@
|
||||
name: flow_matching_generative_finetuning
|
||||
|
||||
init_from_nemo_model: null
|
||||
init_strict: false
|
||||
|
||||
model:
|
||||
type: flow_matching
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
p_cond: 0.9 # Proability of feeding the conditional input into the model.
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
max_utts_evaluation_metrics: 500
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
audio_duration: 6.14 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 768
|
||||
random_offset: true
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
log_config:
|
||||
log_tensorboard: true
|
||||
log_wandb: false
|
||||
max_utts: 8
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.transformerunet.SpectrogramTransformerUNet
|
||||
in_channels: 2 # concatenation of single-channel perturbed and noisy
|
||||
out_channels: 1 # single-channel score estimate
|
||||
depth: 24
|
||||
ff_dropout: 0.1
|
||||
time_hidden_dim: 1024
|
||||
|
||||
flow:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.OptimalTransportFlow
|
||||
sigma_start: 1.0
|
||||
sigma_end: 1e-4
|
||||
|
||||
sampler:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.ConditionalFlowMatchingEulerSampler
|
||||
num_steps: 20
|
||||
time_min: 1e-8
|
||||
time_max: 1.0
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss
|
||||
ndim: 4 # loss is calculated on the score in the encoded domain (batch, channel, dimension, time)
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: 5000
|
||||
warmup_ratio: null
|
||||
min_lr: 0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: 0.2
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_pesq
|
||||
mode: max
|
||||
save_top_k: 3
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: test
|
||||
project: gense
|
||||
@@ -0,0 +1,173 @@
|
||||
name: flow_matching_generative_ssl_pretraining
|
||||
|
||||
model:
|
||||
type: flow_matching
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: true
|
||||
num_outputs: 1
|
||||
p_cond: 0.9 # Proability of feeding the conditional input into the model.
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
max_utts_evaluation_metrics: 125
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
train_ds:
|
||||
shar_path: ???
|
||||
use_lhotse: true
|
||||
truncate_duration: 4.09 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 512
|
||||
truncate_offset_type: random
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: clean_filepath
|
||||
target_key: clean_filepath
|
||||
random_offset: false
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
log_config:
|
||||
log_tensorboard: true
|
||||
log_wandb: false
|
||||
max_utts: 8
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.transformerunet.SpectrogramTransformerUNet
|
||||
in_channels: 2 # concatenation of single-channel perturbed and noisy
|
||||
out_channels: 1 # single-channel score estimate
|
||||
depth: 24
|
||||
ff_dropout: 0.1
|
||||
time_hidden_dim: 1024
|
||||
|
||||
flow:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.OptimalTransportFlow
|
||||
sigma_start: 1.0
|
||||
sigma_end: 1e-4
|
||||
|
||||
sampler:
|
||||
_target_: nemo.collections.audio.parts.submodules.flow.ConditionalFlowMatchingEulerSampler
|
||||
num_steps: 20
|
||||
time_min: 1e-8
|
||||
time_max: 1.0
|
||||
estimator_target: conditional_vector_field # or data
|
||||
|
||||
ssl_pretrain_masking:
|
||||
_target_: nemo.collections.audio.modules.ssl_pretrain_masking.SSLPretrainWithMaskedPatch
|
||||
patch_size: 10
|
||||
mask_fraction: 0.7
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss
|
||||
ndim: 4 # loss is calculated on the score in the encoded domain (batch, channel, dimension, time)
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 5e-5
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: 5000
|
||||
warmup_ratio: null
|
||||
min_lr: 1e-5
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: 10000 # needs to be set for shar datasets
|
||||
limit_train_batches: 1000 # number of batches to train on in each pseudo-epoch
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
use_distributed_sampler: false # required for lhotse
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: 0.2
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_pesq
|
||||
mode: max
|
||||
save_top_k: 3
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,123 @@
|
||||
name: "masking"
|
||||
|
||||
model:
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
audio_duration: 4.0 # in seconds, audio segment duration for training
|
||||
random_offset: true # if the file is longer than audio_duration, use random offset to select a subsegment
|
||||
min_duration: ${model.train_ds.audio_duration}
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
test_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: audio_filepath # key of the input signal path in the manifest
|
||||
target_key: target_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
batch_size: 1 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
mask_estimator:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskEstimatorRNN
|
||||
num_outputs: ${model.num_outputs}
|
||||
num_subbands: 257 # Number of subbands of the input spectrogram
|
||||
num_features: 256 # Number of features at RNN input
|
||||
num_layers: 5 # Number of RNN layers
|
||||
bidirectional: true # Use bi-directional RNN
|
||||
|
||||
mask_processor:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskReferenceChannel # Apply mask on the reference channel
|
||||
ref_channel: 0 # Reference channel for the output
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.SDRLoss
|
||||
scale_invariant: true # Use scale-invariant SDR
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sdr: # output SDR
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
test:
|
||||
sdr_ch0: # SDR on output channel 0
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
channel: 0
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: False # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
create_tensorboard_logger: true
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: "val_loss"
|
||||
mode: "min"
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,119 @@
|
||||
name: "masking_with_online_augmenatation"
|
||||
|
||||
model:
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
|
||||
train_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with speech signals for augmentation (including custom "target_recording" field with the same signals)
|
||||
truncate_duration: 4.0 # Number of STFT time frames = 1 + truncate_duration // encoder.hop_length = 256
|
||||
truncate_offset_type: random # if the file is longer than truncate_duration, use random offset to select a subsegment
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
rir_enabled: true # enable room impulse response augmentation
|
||||
rir_path: ??? # path to Lhotse recordings manifest with room impulse response signals
|
||||
noise_path: ??? # path to Lhotse cuts manifest with noise signals
|
||||
|
||||
validation_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with noisy speech signals (including custom "target_recording" field with the clean signals)
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
test_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with noisy speech signals (including custom "target_recording" field with the clean signals)
|
||||
batch_size: 1 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: 512 # Length of the window and FFT for calculating spectrogram
|
||||
hop_length: 256 # Hop length for calculating spectrogram
|
||||
|
||||
mask_estimator:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskEstimatorRNN
|
||||
num_outputs: ${model.num_outputs}
|
||||
num_subbands: 257 # Number of subbands of the input spectrogram
|
||||
num_features: 256 # Number of features at RNN input
|
||||
num_layers: 5 # Number of RNN layers
|
||||
bidirectional: true # Use bi-directional RNN
|
||||
|
||||
mask_processor:
|
||||
_target_: nemo.collections.audio.modules.masking.MaskReferenceChannel # Apply mask on the reference channel
|
||||
ref_channel: 0 # Reference channel for the output
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.SDRLoss
|
||||
scale_invariant: true # Use scale-invariant SDR
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sdr: # output SDR
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
test:
|
||||
sdr_ch0: # SDR on output channel 0
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
channel: 0
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: False # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
create_tensorboard_logger: true
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: "val_loss"
|
||||
mode: "min"
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,122 @@
|
||||
name: "BNR2"
|
||||
|
||||
model:
|
||||
type: bnr
|
||||
sample_rate: 16000
|
||||
fft_length: 1920
|
||||
hop_length: 480
|
||||
num_mels: 320
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
segment: 4
|
||||
|
||||
train: # Parameters related to training
|
||||
enable_weight_norm: true
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 0.0005
|
||||
sched:
|
||||
name: StepLR
|
||||
gamma: 0.999
|
||||
step_size: 2
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath # key of the input signal path in the manifest
|
||||
target_key: speech_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
audio_duration: 4.0 # in seconds, audio segment duration for training
|
||||
random_offset: true # if the file is longer than audio_duration, use random offset to select a subsegment
|
||||
min_duration: ${model.train_ds.audio_duration}
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath # key of the input signal path in the manifest
|
||||
target_key: speech_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
audio_duration: 10.0 # in seconds, audio segment duration for validation
|
||||
min_duration: ${model.validation_ds.audio_duration}
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
test_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath # key of the input signal path in the manifest
|
||||
target_key: speech_filepath # key of the target signal path in the manifest
|
||||
target_channel_selector: 0 # target signal is the first channel from files in target_key
|
||||
audio_duration: 10.0 # in seconds, audio segment duration for validation
|
||||
min_duration: ${model.test_ds.audio_duration}
|
||||
batch_size: 1 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.maxine.CombinedLoss
|
||||
sample_rate: ${model.sample_rate}
|
||||
fft_length: ${model.fft_length}
|
||||
hop_length: ${model.hop_length}
|
||||
num_mels: ${model.num_mels}
|
||||
sisnr_loss_weight: 1
|
||||
spectral_loss_weight: 15
|
||||
asr_loss_weight: 1
|
||||
use_asr_loss: true
|
||||
use_mel_spec: true
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sdr: # output SDR
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
test:
|
||||
sdr_ch0: # SDR on output channel 0
|
||||
_target_: torchmetrics.audio.SignalDistortionRatio
|
||||
channel: 0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: gpu
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: 5
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: False # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
create_tensorboard_logger: true
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: "val_loss"
|
||||
mode: "min"
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,130 @@
|
||||
name: "predictive_model"
|
||||
|
||||
model:
|
||||
type: predictive
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
audio_duration: 2.04 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 256
|
||||
random_offset: true
|
||||
normalization_signal: input_signal
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.ncsnpp.SpectrogramNoiseConditionalScoreNetworkPlusPlus
|
||||
in_channels: 1 # single-channel noisy input
|
||||
out_channels: 1 # single-channel estimate
|
||||
num_res_blocks: 3 # increased number of res blocks
|
||||
pad_time_to: 64 # pad to 64 frames for the time dimension
|
||||
pad_dimension_to: 0 # no padding in the frequency dimension
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss # computed in the time domain
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: False # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,141 @@
|
||||
name: predictive_conformer
|
||||
|
||||
model:
|
||||
type: predictive
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
# non-streaming config, use input normalization
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
|
||||
train_ds:
|
||||
shar_path: ???
|
||||
use_lhotse: true
|
||||
truncate_duration: 4.09 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 512
|
||||
truncate_offset_type: random
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.conformer.SpectrogramConformer
|
||||
in_channels: 1 # single-channel noisy input
|
||||
out_channels: 1 # single-channel estimate
|
||||
feat_in: 256 # input feature dimension = number of subbands
|
||||
n_layers: 8 # number of layers in the model
|
||||
d_model: 512 # the hidden size of the model
|
||||
subsampling_factor: 1 # subsampling factor for the model
|
||||
self_attention_model: 'rel_pos'
|
||||
n_heads: 8 # number of heads for the model
|
||||
# streaming-related arguments
|
||||
# - this is a non-streaming config
|
||||
conv_context_size: null
|
||||
conv_norm_type: 'layer_norm'
|
||||
causal_downsampling: False
|
||||
att_context_size: [-1, -1]
|
||||
att_context_style: 'regular'
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss # computed in the time domain
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-3
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: null
|
||||
warmup_ratio: 0.1
|
||||
min_lr: 1e-5
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: ??? # needs to be set for shar datasets
|
||||
limit_train_batches: ??? # number of batches to train on in each pseudo-epoch
|
||||
val_check_interval: ??? # run validation after this many training steps
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
use_distributed_sampler: false # required for lhotse
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 100 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: null # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,147 @@
|
||||
name: predictive_conformer_unet
|
||||
|
||||
model:
|
||||
type: predictive
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
# non-streaming config, use input normalization
|
||||
normalize_input: true # normalize the input signal to 0dBFS
|
||||
|
||||
train_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with speech signals for augmentation (including custom "target_recording" field with the same signals)
|
||||
truncate_duration: 2.04 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 256
|
||||
truncate_offset_type: random # if the file is longer than truncate_duration, use random offset to select a subsegment
|
||||
batch_size: 32 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with noisy speech signals (including custom "target_recording" field with the clean signals)
|
||||
batch_size: 4 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.conformer_unet.SpectrogramConformerUNet
|
||||
in_channels: 1 # single-channel noisy input
|
||||
out_channels: 1 # single-channel estimate
|
||||
feat_in: 256 # input feature dimension = number of subbands
|
||||
n_layers: 8 # number of layers in the model
|
||||
d_model: 512 # the hidden size of the model
|
||||
subsampling_factor: 1 # subsampling factor for the model
|
||||
self_attention_model: 'rel_pos'
|
||||
n_heads: 8 # number of heads for the model
|
||||
# streaming-related arguments
|
||||
# - this is a non-streaming config
|
||||
conv_context_size: null
|
||||
conv_norm_type: 'layer_norm'
|
||||
causal_downsampling: False
|
||||
att_context_size: [-1, -1]
|
||||
att_context_style: 'regular'
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss # computed in the time domain
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 0.0
|
||||
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # run validation after this many training steps
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 100 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
use_distributed_sampler: false # required for lhotse
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,164 @@
|
||||
name: schroedinger_bridge
|
||||
|
||||
model:
|
||||
type: schroedinger_bridge
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
normalize_input: true
|
||||
max_utts_evaluation_metrics: 50 # metric calculation needs full inference and is slow, so we limit to first few files
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
audio_duration: 2.04 # 256 frames
|
||||
random_offset: true
|
||||
normalize_input: ${model.normalize_input}
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
normalize_input: false # load data as is for validation, the model will normalize it for inference
|
||||
batch_size: 4
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.ncsnpp.SpectrogramNoiseConditionalScoreNetworkPlusPlus
|
||||
in_channels: 2 # concatenation of single-channel perturbed and noisy
|
||||
out_channels: 1 # single-channel estimate
|
||||
conditioned_on_time: true
|
||||
num_res_blocks: 3 # increased number of res blocks
|
||||
pad_time_to: 64 # pad to 64 frames for the time dimension
|
||||
pad_dimension_to: 0 # no padding in the frequency dimension
|
||||
|
||||
estimator_output: data_prediction
|
||||
|
||||
noise_schedule:
|
||||
_target_: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBNoiseScheduleVE
|
||||
k: 2.6
|
||||
c: 0.4
|
||||
time_min: 1e-4
|
||||
time_max: 1.0
|
||||
num_steps: 1000 # num steps for the forward process
|
||||
|
||||
sampler:
|
||||
_target_: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBSampler
|
||||
time_min: 1e-4
|
||||
time_max: 1.0
|
||||
num_steps: 50 # num steps for the reverse process
|
||||
|
||||
# Loss in the encoded domain
|
||||
loss_encoded:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss
|
||||
ndim: 4 # loss is calculated on the score in the encoded domain (batch, channel, dimension, time)
|
||||
|
||||
# Loss in the time domain
|
||||
loss_time:
|
||||
_target_: nemo.collections.audio.losses.audio.MAELoss
|
||||
loss_time_weight: 0.001
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 5 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_pesq
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,149 @@
|
||||
name: score_based_generative_model
|
||||
|
||||
model:
|
||||
type: score_based
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
normalize_input: true
|
||||
max_utts_evaluation_metrics: 50 # metric calculation needs full inference and is slow, so we limit to first few files
|
||||
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
audio_duration: 2.04 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 256
|
||||
random_offset: true
|
||||
normalization_signal: input_signal
|
||||
batch_size: 8 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
normalize_input: false # load data as is for validation, the model will normalize it for inference
|
||||
batch_size: 4
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.ncsnpp.SpectrogramNoiseConditionalScoreNetworkPlusPlus
|
||||
in_channels: 2 # concatenation of single-channel perturbed and noisy
|
||||
out_channels: 1 # single-channel score estimate
|
||||
conditioned_on_time: true
|
||||
num_res_blocks: 3 # increased number of res blocks
|
||||
pad_time_to: 64 # pad to 64 frames for the time dimension
|
||||
pad_dimension_to: 0 # no padding in the frequency dimension
|
||||
|
||||
sde:
|
||||
_target_: nemo.collections.audio.parts.submodules.diffusion.OrnsteinUhlenbeckVarianceExplodingSDE
|
||||
stiffness: 1.5
|
||||
std_min: 0.05
|
||||
std_max: 0.5
|
||||
num_steps: 1000
|
||||
|
||||
sampler:
|
||||
_target_: nemo.collections.audio.parts.submodules.diffusion.PredictorCorrectorSampler
|
||||
predictor: reverse_diffusion
|
||||
corrector: annealed_langevin_dynamics
|
||||
num_steps: 50
|
||||
num_corrector_steps: 1
|
||||
snr: 0.5
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss
|
||||
ndim: 4 # loss is calculated on the score in the encoded domain (batch, channel, dimension, time)
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.999]
|
||||
weight_decay: 0.0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 25 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,141 @@
|
||||
name: streaming_predictive_conformer
|
||||
|
||||
model:
|
||||
type: predictive
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
# streaming config, without input normalization
|
||||
normalize_input: false
|
||||
|
||||
train_ds:
|
||||
shar_path: ???
|
||||
use_lhotse: true
|
||||
truncate_duration: 4.09 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 512
|
||||
truncate_offset_type: random
|
||||
batch_size: 64 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
input_key: noisy_filepath
|
||||
target_key: clean_filepath
|
||||
batch_size: 8
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.conformer.SpectrogramConformer
|
||||
in_channels: 1 # single-channel noisy input
|
||||
out_channels: 1 # single-channel estimate
|
||||
feat_in: 256 # input feature dimension = number of subbands
|
||||
n_layers: 8 # number of layers in the model
|
||||
d_model: 512 # the hidden size of the model
|
||||
subsampling_factor: 1 # subsampling factor for the model
|
||||
self_attention_model: 'rel_pos'
|
||||
n_heads: 8 # number of heads for the model
|
||||
# streaming-related arguments
|
||||
# - streaming config with causal convolutions and limited attention context
|
||||
conv_context_size: 'causal'
|
||||
conv_norm_type: 'layer_norm'
|
||||
causal_downsampling: True
|
||||
att_context_size: [102, 16]
|
||||
att_context_style: 'chunked_limited'
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss # computed in the time domain
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
|
||||
optim:
|
||||
name: adamw
|
||||
lr: 1e-3
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 1e-3
|
||||
|
||||
# scheduler setup
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
# scheduler config override
|
||||
warmup_steps: null
|
||||
warmup_ratio: 0.1
|
||||
min_lr: 1e-5
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: ??? # needs to be set for shar datasets
|
||||
limit_train_batches: ??? # number of batches to train on in each pseudo-epoch
|
||||
val_check_interval: ??? # run validation after this many training steps
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
use_distributed_sampler: false # required for lhotse
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 100 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: null # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,146 @@
|
||||
name: streaming_predictive_conformer_unet
|
||||
|
||||
model:
|
||||
type: predictive
|
||||
sample_rate: 16000
|
||||
skip_nan_grad: false
|
||||
num_outputs: 1
|
||||
# streaming config, without input normalization
|
||||
normalize_input: false
|
||||
|
||||
train_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with speech signals for augmentation (including custom "target_recording" field with the same signals)
|
||||
truncate_duration: 2.04 # Number of STFT time frames = 1 + audio_duration // encoder.hop_length = 256
|
||||
truncate_offset_type: random # if the file is longer than truncate_duration, use random offset to select a subsegment
|
||||
batch_size: 32 # batch size may be increased based on the available memory
|
||||
shuffle: true
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
use_lhotse: true # enable Lhotse data loader
|
||||
cuts_path: ??? # path to Lhotse cuts manifest with noisy speech signals (including custom "target_recording" field with the clean signals)
|
||||
batch_size: 4 # batch size may be increased based on the available memory
|
||||
shuffle: false
|
||||
num_workers: 4
|
||||
pin_memory: true
|
||||
|
||||
encoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
|
||||
fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
|
||||
hop_length: 128
|
||||
magnitude_power: 0.5
|
||||
scale: 0.33
|
||||
|
||||
decoder:
|
||||
_target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
|
||||
fft_length: ${model.encoder.fft_length}
|
||||
hop_length: ${model.encoder.hop_length}
|
||||
magnitude_power: ${model.encoder.magnitude_power}
|
||||
scale: ${model.encoder.scale}
|
||||
|
||||
estimator:
|
||||
_target_: nemo.collections.audio.parts.submodules.conformer_unet.SpectrogramConformerUNet
|
||||
in_channels: 1 # single-channel noisy input
|
||||
out_channels: 1 # single-channel estimate
|
||||
feat_in: 256 # input feature dimension = number of subbands
|
||||
n_layers: 8 # number of layers in the model
|
||||
d_model: 512 # the hidden size of the model
|
||||
subsampling_factor: 1 # subsampling factor for the model
|
||||
self_attention_model: 'rel_pos'
|
||||
n_heads: 8 # number of heads for the model
|
||||
# streaming-related arguments
|
||||
# - streaming config with causal convolutions and limited attention context
|
||||
conv_context_size: 'causal'
|
||||
conv_norm_type: 'layer_norm'
|
||||
causal_downsampling: True
|
||||
att_context_size: [102, 16]
|
||||
att_context_style: 'chunked_limited'
|
||||
|
||||
loss:
|
||||
_target_: nemo.collections.audio.losses.audio.MSELoss # computed in the time domain
|
||||
|
||||
metrics:
|
||||
val:
|
||||
sisdr: # output SI-SDR
|
||||
_target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
|
||||
estoi: # output ESTOI
|
||||
_target_: torchmetrics.audio.ShortTimeObjectiveIntelligibility
|
||||
fs: ${model.sample_rate}
|
||||
extended: true
|
||||
pesq: # output PESQ
|
||||
_target_: torchmetrics.audio.PerceptualEvaluationSpeechQuality
|
||||
fs: ${model.sample_rate}
|
||||
mode: wb
|
||||
|
||||
optim:
|
||||
name: adam
|
||||
lr: 1e-4
|
||||
# optimizer arguments
|
||||
betas: [0.9, 0.98]
|
||||
weight_decay: 0.0
|
||||
|
||||
trainer:
|
||||
devices: -1 # number of GPUs, -1 would use all available GPUs
|
||||
num_nodes: 1
|
||||
max_epochs: -1
|
||||
max_steps: -1 # computed at runtime if not set
|
||||
val_check_interval: 1.0 # run validation after this many training steps
|
||||
accelerator: auto
|
||||
strategy: ddp
|
||||
accumulate_grad_batches: 1
|
||||
gradient_clip_val: null
|
||||
precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
|
||||
log_every_n_steps: 100 # Interval of logging.
|
||||
enable_progress_bar: true
|
||||
num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
|
||||
check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
|
||||
sync_batchnorm: true
|
||||
enable_checkpointing: false # Provided by exp_manager
|
||||
logger: false # Provided by exp_manager
|
||||
use_distributed_sampler: false # required for lhotse
|
||||
|
||||
exp_manager:
|
||||
exp_dir: null
|
||||
name: ${name}
|
||||
# use exponential moving average for model parameters
|
||||
ema:
|
||||
enable: true
|
||||
decay: 0.999 # decay rate
|
||||
cpu_offload: false # offload EMA parameters to CPU to save GPU memory
|
||||
every_n_steps: 1 # how often to update EMA weights
|
||||
validate_original_weights: false # use original weights for validation calculation?
|
||||
|
||||
# logging
|
||||
create_tensorboard_logger: true
|
||||
|
||||
# checkpointing
|
||||
create_checkpoint_callback: true
|
||||
checkpoint_callback_params:
|
||||
# in case of multiple validation sets, first one is used
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
save_top_k: 5
|
||||
always_save_nemo: true # saves the checkpoints as nemo files instead of PTL checkpoints
|
||||
|
||||
# early stopping
|
||||
create_early_stopping_callback: true
|
||||
early_stopping_callback_params:
|
||||
monitor: val_sisdr
|
||||
mode: max
|
||||
min_delta: 0.0
|
||||
patience: 20 # patience in terms of check_val_every_n_epoch
|
||||
verbose: true
|
||||
strict: false # Should be False to avoid a runtime error where EarlyStopping says monitor is unavailable, which sometimes happens with resumed training.
|
||||
|
||||
resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
|
||||
# you need to set these two to true to continue the training
|
||||
resume_if_exists: false
|
||||
resume_ignore_no_checkpoint: false
|
||||
|
||||
# You may use this section to create a W&B logger
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
name: null
|
||||
project: null
|
||||
@@ -0,0 +1,277 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import contextlib
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field, is_dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
import lightning.pytorch as pl
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.audio.models.audio_to_audio import AudioToAudioModel
|
||||
from nemo.core.config import hydra_runner
|
||||
from nemo.utils import logging, model_utils
|
||||
|
||||
|
||||
"""
|
||||
Process audio file on a single CPU/GPU. Useful for processing of moderate amounts of audio data.
|
||||
|
||||
# Arguments
|
||||
model_path: path to .nemo checkpoint for an AudioToAudioModel
|
||||
pretrained_name: name of a pretrained AudioToAudioModel model (from NGC registry)
|
||||
audio_dir: path to directory with audio files
|
||||
dataset_manifest: path to dataset JSON manifest file (in NeMo format)
|
||||
max_utts: maximum number of utterances to process
|
||||
|
||||
input_channel_selector: list of channels to take from audio files, defaults to `None` and takes all available channels
|
||||
input_key: key for audio filepath in the manifest file, defaults to `audio_filepath`
|
||||
|
||||
output_dir: Directory where processed files will be saved
|
||||
output_filename: Output filename where manifest pointing to processed files will be written
|
||||
batch_size: batch size during inference
|
||||
|
||||
cuda: Optional int to enable or disable execution of model on certain CUDA device.
|
||||
amp: Bool to decide if Automatic Mixed Precision should be used during inference
|
||||
audio_type: Str filetype of the audio. Supported = wav, flac, mp3
|
||||
|
||||
overwrite_output: Bool which when set allowes repeated processing runs to overwrite previous results.
|
||||
|
||||
# Usage
|
||||
AudioToAudioModel can be specified by either `model_path` or `pretrained_name`.
|
||||
Data for processing can be defined with either `audio_dir` or `dataset_manifest`.
|
||||
Processed audio is saved in `output_dir`, and a manifest for processed files is saved
|
||||
in `output_filename`.
|
||||
|
||||
```
|
||||
python process_audio.py \
|
||||
model_path=null \
|
||||
pretrained_name=null \
|
||||
audio_dir="" \
|
||||
dataset_manifest="" \
|
||||
input_channel_selector=[] \
|
||||
output_dir="" \
|
||||
output_filename="" \
|
||||
batch_size=1 \
|
||||
cuda=0 \
|
||||
amp=True
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessConfig:
|
||||
# Required configs
|
||||
model_path: Optional[str] = None # Path to a .nemo file
|
||||
pretrained_name: Optional[str] = None # Name of a pretrained model
|
||||
audio_dir: Optional[str] = None # Path to a directory which contains audio files
|
||||
dataset_manifest: Optional[str] = None # Path to dataset's JSON manifest
|
||||
max_utts: Optional[int] = None # max number of utterances to process
|
||||
|
||||
# Audio configs
|
||||
input_channel_selector: Optional[List] = None # Union types not supported Optional[Union[List, int]]
|
||||
input_key: Optional[str] = None # Can be used with a manifest
|
||||
|
||||
# General configs
|
||||
output_dir: Optional[str] = None
|
||||
output_filename: Optional[str] = None
|
||||
batch_size: int = 1
|
||||
num_workers: int = 0
|
||||
|
||||
# Override model config
|
||||
override_config_path: Optional[str] = None # path to a yaml config that will override the internal config file
|
||||
|
||||
# Override sampler config
|
||||
# For example, to set number of steps, use `++sampler.num_samples=42`
|
||||
sampler: dict = field(default_factory=dict)
|
||||
|
||||
# Set `cuda` to int to define CUDA device. If 'None', will look for CUDA
|
||||
# device anyway, and do inference on CPU only if CUDA device is not found.
|
||||
# If `cuda` is a negative number, inference will be on CPU only.
|
||||
cuda: Optional[int] = None
|
||||
amp: bool = False
|
||||
audio_type: str = "wav"
|
||||
|
||||
# Recompute model predictions, even if the output folder exists.
|
||||
overwrite_output: bool = False
|
||||
|
||||
|
||||
@hydra_runner(config_name="ProcessConfig", schema=ProcessConfig)
|
||||
def main(cfg: ProcessConfig) -> ProcessConfig:
|
||||
logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
|
||||
|
||||
if is_dataclass(cfg):
|
||||
cfg = OmegaConf.structured(cfg)
|
||||
|
||||
if cfg.model_path is None and cfg.pretrained_name is None:
|
||||
raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None!")
|
||||
if cfg.audio_dir is None and cfg.dataset_manifest is None:
|
||||
raise ValueError("Both cfg.audio_dir and cfg.dataset_manifest cannot be None!")
|
||||
|
||||
# setup GPU
|
||||
if cfg.cuda is None:
|
||||
if torch.cuda.is_available():
|
||||
device = [0] # use 0th CUDA device
|
||||
accelerator = 'gpu'
|
||||
else:
|
||||
device = 1
|
||||
accelerator = 'cpu'
|
||||
else:
|
||||
device = [cfg.cuda]
|
||||
accelerator = 'gpu'
|
||||
|
||||
map_location = torch.device('cuda:{}'.format(device[0]) if accelerator == 'gpu' else 'cpu')
|
||||
|
||||
# setup model
|
||||
if cfg.model_path is not None:
|
||||
# restore model from .nemo file path
|
||||
model_cfg = AudioToAudioModel.restore_from(restore_path=cfg.model_path, return_config=True)
|
||||
classpath = model_cfg.target # original class path
|
||||
imported_class = model_utils.import_class_by_path(classpath) # type: AudioToAudioModel
|
||||
logging.info(f"Restoring model : {imported_class.__name__}")
|
||||
audio_to_audio_model = imported_class.restore_from(
|
||||
restore_path=cfg.model_path, override_config_path=cfg.override_config_path, map_location=map_location
|
||||
) # type: AudioToAudioModel
|
||||
model_name = os.path.splitext(os.path.basename(cfg.model_path))[0]
|
||||
else:
|
||||
# restore model by name
|
||||
audio_to_audio_model = AudioToAudioModel.from_pretrained(
|
||||
model_name=cfg.pretrained_name, map_location=map_location
|
||||
) # type: AudioToAudioModel
|
||||
model_name = cfg.pretrained_name
|
||||
|
||||
trainer = pl.Trainer(devices=device, accelerator=accelerator)
|
||||
audio_to_audio_model.set_trainer(trainer)
|
||||
audio_to_audio_model = audio_to_audio_model.eval()
|
||||
|
||||
# override sampler if necessary
|
||||
if cfg.sampler:
|
||||
logging.info('Overriding sampler with %s', cfg.sampler)
|
||||
|
||||
if hasattr(audio_to_audio_model, 'sampler'):
|
||||
for key, value in cfg.sampler.items():
|
||||
if not hasattr(audio_to_audio_model.sampler, key):
|
||||
raise RuntimeError(f'Model sampler does not have attribute {key}')
|
||||
logging.debug('Try to set model.sampler.%s to %s', key, value)
|
||||
setattr(audio_to_audio_model.sampler, key, value)
|
||||
if getattr(audio_to_audio_model.sampler, key) != value:
|
||||
raise RuntimeError(f'Failed to set model sampler attribute {key} to {value}')
|
||||
logging.info('model.sampler.%s was set to %s', key, value)
|
||||
else:
|
||||
raise RuntimeError('Model does not have a sampler')
|
||||
|
||||
if cfg.audio_dir is not None:
|
||||
input_dir = cfg.audio_dir
|
||||
filepaths = list(glob.glob(os.path.join(cfg.audio_dir, f"**/*.{cfg.audio_type}"), recursive=True))
|
||||
else:
|
||||
# get filenames from manifest
|
||||
filepaths = []
|
||||
if os.stat(cfg.dataset_manifest).st_size == 0:
|
||||
raise RuntimeError(f"The input dataset_manifest {cfg.dataset_manifest} is empty.")
|
||||
|
||||
input_key = 'audio_filepath' if cfg.input_key is None else cfg.input_key
|
||||
manifest_dir = Path(cfg.dataset_manifest).parent
|
||||
with open(cfg.dataset_manifest, 'r') as f:
|
||||
for line in f:
|
||||
item = json.loads(line)
|
||||
audio_file = Path(item[input_key])
|
||||
if not audio_file.is_file() and not audio_file.is_absolute():
|
||||
audio_file = manifest_dir / audio_file
|
||||
filepaths.append(str(audio_file.absolute()))
|
||||
|
||||
# common path for all files
|
||||
common_path = os.path.commonpath(filepaths)
|
||||
if Path(common_path).is_relative_to(manifest_dir):
|
||||
# if all paths are relative to the manifest, use manifest dir as input dir
|
||||
input_dir = manifest_dir
|
||||
else:
|
||||
# use the parent of the common path as input dir
|
||||
input_dir = Path(common_path).parent
|
||||
|
||||
if cfg.max_utts is not None:
|
||||
# Limit the number of utterances to process
|
||||
filepaths = filepaths[: cfg.max_utts]
|
||||
|
||||
logging.info(f"\nProcessing {len(filepaths)} files...\n")
|
||||
|
||||
# setup AMP (optional)
|
||||
if cfg.amp and torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'):
|
||||
logging.info("AMP enabled!\n")
|
||||
autocast = torch.cuda.amp.autocast
|
||||
else:
|
||||
|
||||
@contextlib.contextmanager
|
||||
def autocast():
|
||||
yield
|
||||
|
||||
# Compute output filename
|
||||
if cfg.output_dir is None:
|
||||
# create default output filename
|
||||
if cfg.audio_dir is not None:
|
||||
cfg.output_dir = os.path.dirname(os.path.join(cfg.audio_dir, '.')) + f'_processed_{model_name}'
|
||||
else:
|
||||
cfg.output_dir = os.path.dirname(cfg.dataset_manifest) + f'_processed_{model_name}'
|
||||
|
||||
# Compute output filename
|
||||
if cfg.output_filename is None:
|
||||
# create default output filename
|
||||
cfg.output_filename = cfg.output_dir.rstrip('/') + '_manifest.json'
|
||||
|
||||
# if transcripts should not be overwritten, and already exists, skip re-transcription step and return
|
||||
if not cfg.overwrite_output and os.path.exists(cfg.output_dir):
|
||||
raise RuntimeError(
|
||||
f"Previous output found at {cfg.output_dir}, and flag `overwrite_output`"
|
||||
f"is {cfg.overwrite_output}. Returning without processing."
|
||||
)
|
||||
|
||||
# Process audio
|
||||
with autocast():
|
||||
with torch.no_grad():
|
||||
paths2processed_files = audio_to_audio_model.process(
|
||||
paths2audio_files=filepaths,
|
||||
output_dir=cfg.output_dir,
|
||||
batch_size=cfg.batch_size,
|
||||
num_workers=cfg.num_workers,
|
||||
input_channel_selector=cfg.input_channel_selector,
|
||||
input_dir=input_dir,
|
||||
)
|
||||
|
||||
logging.info(f"Finished processing {len(filepaths)} files!")
|
||||
logging.info(f"Processed audio is available in the output directory: {cfg.output_dir}")
|
||||
|
||||
# Prepare new/updated manifest with a new key for processed audio
|
||||
with open(cfg.output_filename, 'w', encoding='utf-8') as f:
|
||||
if cfg.dataset_manifest is not None:
|
||||
with open(cfg.dataset_manifest, 'r') as fr:
|
||||
for idx, line in enumerate(fr):
|
||||
item = json.loads(line)
|
||||
item['processed_audio_filepath'] = paths2processed_files[idx]
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
if cfg.max_utts is not None and idx >= cfg.max_utts - 1:
|
||||
break
|
||||
else:
|
||||
for idx, processed_file in enumerate(paths2processed_files):
|
||||
item = {'processed_audio_filepath': processed_file}
|
||||
f.write(json.dumps(item) + "\n")
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
@@ -0,0 +1,193 @@
|
||||
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from itertools import islice
|
||||
from pathlib import Path
|
||||
|
||||
import hydra
|
||||
import lhotse
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from lhotse import CutSet, MonoCut, Recording
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
from tqdm import tqdm
|
||||
|
||||
from nemo.collections.audio.data.audio_to_audio_lhotse import LhotseAudioToTargetDataset
|
||||
from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig, get_lhotse_dataloader_from_config
|
||||
|
||||
"""
|
||||
The purpose of this script is to save online-augmented data as provided by NeMo Lhotse dataloader.
|
||||
The script piggybacks on a train_ds section of an existing model configuration file.
|
||||
|
||||
Intended use cases are: 1) preparing a validation set, 2) debugging.
|
||||
|
||||
Usage example:
|
||||
$ python examples/audio/save_augmented.py \
|
||||
+input_cuts=some_path/cuts.jsonl \
|
||||
+output_cuts=some_other_path/cuts.gsm_and_clipping_augmented.jsonl \
|
||||
+keep_directory_structure=true \
|
||||
model.sample_rate=48000 \
|
||||
++model.train_ds.rir_enabled=true \
|
||||
++model.train_ds.rir_path=path/to/rir_manifest.jsonl
|
||||
|
||||
Assumptions:
|
||||
- input data are described as a Lhotse CutSet in a JSONL file
|
||||
- consists of simple MonoCuts with Recording paths relative to the Cuts manifest
|
||||
- the parent directory of the output cuts must exist
|
||||
|
||||
Requires additional config parameters `input_cuts` and `output_cuts`.
|
||||
Produces:
|
||||
- %output_cuts_parent_dir%/audio/
|
||||
- %output_cuts_parent_dir%/%output_cuts_filename%.jsonl
|
||||
where the audio folder contains the augmented and clean signals, respectively, with `.input.flac` and `.output.flac` suffixes.
|
||||
|
||||
If `keep_directory_structure` provided and is True, the script will preserve the directory structure of the input cuts.
|
||||
|
||||
Text is preserved from the input cuts if possible.
|
||||
|
||||
Optional config parameter `num_samples` can be used to limit the number of samples to save (but not more than input dataloader size).
|
||||
If not specified, the dataloader is used until exhausted.
|
||||
"""
|
||||
|
||||
|
||||
def check_input_cuts(input_cuts_path: Path) -> None:
|
||||
"""Validate that input cuts are well-formed MonoCuts with relative recording paths that exist on disk."""
|
||||
assert input_cuts_path.exists(), "input_cuts must exist"
|
||||
assert input_cuts_path.suffix == '.jsonl', "input_cuts must be a .jsonl file"
|
||||
assert input_cuts_path.parent.exists(), "input_cuts parent directory must exist"
|
||||
cuts = lhotse.CutSet.from_file(input_cuts_path)
|
||||
for i, cut in enumerate(cuts):
|
||||
assert isinstance(cut, MonoCut), f"{i}th cut is a {type(cut)}, not a MonoCut"
|
||||
assert len(cut.recording.sources) == 1, f"{i}th cut has {len(cut.recording.sources)} sources"
|
||||
assert cut.recording.sources[0].source is not None, f"{i}th cut has no audio source specified"
|
||||
|
||||
recording_path = Path(cut.recording.sources[0].source)
|
||||
assert not recording_path.is_absolute(), f"{i}th cut's recording source is an absolute path: {recording_path}"
|
||||
|
||||
recording_path_full = input_cuts_path.parent / recording_path
|
||||
assert recording_path_full.exists(), f"{i}th cut's recording source file does not exist: {recording_path_full}"
|
||||
|
||||
|
||||
@hydra.main(config_path="conf", config_name="flow_matching_generative_finetuning.yaml")
|
||||
def main(cfg: DictConfig):
|
||||
assert (
|
||||
cfg.get("input_cuts", None) is not None
|
||||
), "input_cuts is required, please override (for example, +input_cuts=some_path/cuts.jsonl)"
|
||||
assert (
|
||||
cfg.get("output_cuts", None) is not None
|
||||
), "output_cuts is required, please override (for example, +output_cuts=some_path/cuts.augmented.jsonl)"
|
||||
num_samples = cfg.get("num_samples", None)
|
||||
sample_rate = cfg.model.sample_rate
|
||||
keep_directory_structure = cfg.get("keep_directory_structure", False)
|
||||
|
||||
input_cuts_path = Path(cfg.input_cuts)
|
||||
output_cuts_path = Path(cfg.output_cuts)
|
||||
check_input_cuts(input_cuts_path) # throws an exception if they aren't ok
|
||||
|
||||
assert output_cuts_path.parent.exists(), f"output_cuts parent directory must exist: {output_cuts_path.parent}"
|
||||
|
||||
OmegaConf.set_struct(cfg, True)
|
||||
OmegaConf.update(cfg, "model.train_ds.cuts_path", str(input_cuts_path), force_add=True)
|
||||
OmegaConf.update(cfg, "model.train_ds.shuffle", False) # ensure deterministic behavior
|
||||
OmegaConf.update(cfg, "model.train_ds.batch_size", 1)
|
||||
OmegaConf.update(cfg, "model.train_ds.shard_seed", 0, force_add=True) # ensure deterministic behavior
|
||||
if cfg.model.train_ds.get("sample_rate", None) != sample_rate:
|
||||
OmegaConf.update(cfg, "model.train_ds.sample_rate", sample_rate, force_add=True)
|
||||
|
||||
# Disable bucketing to preserve original cut ordering (DynamicBucketingSampler reorders by duration).
|
||||
# Also clear bucket params that would cause _auto_detect_bucketing_and_validate_batch_size to re-enable it.
|
||||
OmegaConf.update(cfg, "model.train_ds.use_bucketing", False, force_add=True)
|
||||
_defaults = LhotseDataLoadingConfig()
|
||||
for key in ("bucket_batch_size", "bucket_duration_bins"):
|
||||
OmegaConf.update(cfg, f"model.train_ds.{key}", getattr(_defaults, key), force_add=True)
|
||||
|
||||
# Reset all filters to pass-through defaults — we want a 1:1 mapping from input to output cuts,
|
||||
# so no cuts should be silently dropped by model-config filter settings.
|
||||
for key in (
|
||||
"min_duration",
|
||||
"max_duration",
|
||||
"min_tps",
|
||||
"max_tps",
|
||||
"min_tokens",
|
||||
"max_tokens",
|
||||
"max_cer",
|
||||
"min_context_speaker_similarity",
|
||||
):
|
||||
OmegaConf.update(cfg, f"model.train_ds.{key}", getattr(_defaults, key), force_add=True)
|
||||
|
||||
dataloader = get_lhotse_dataloader_from_config(
|
||||
OmegaConf.create(cfg.model.train_ds), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()
|
||||
)
|
||||
|
||||
cuts = lhotse.CutSet.from_file(input_cuts_path)
|
||||
if num_samples is None:
|
||||
num_samples = len(cuts)
|
||||
|
||||
with CutSet.open_writer(output_cuts_path) as writer:
|
||||
for i, (sample, original_cut) in enumerate(
|
||||
tqdm(zip(islice(dataloader, num_samples), cuts), total=num_samples)
|
||||
):
|
||||
# batch_size is 1, so we can access the first element
|
||||
input_audio = sample['input_signal'][0].numpy()
|
||||
output_audio = sample['target_signal'][0].numpy()
|
||||
|
||||
# if necessary, apply negative gain to avoid clipping
|
||||
if (coeff := max(np.max(np.abs(input_audio)), np.max(np.abs(output_audio)))) > 1.0:
|
||||
input_audio = input_audio / coeff
|
||||
output_audio = output_audio / coeff
|
||||
|
||||
if keep_directory_structure:
|
||||
# definitely a relative path because we checked for that earlier
|
||||
input_relative_path = Path(original_cut.recording.sources[0].source)
|
||||
|
||||
input_path = output_cuts_path.parent / input_relative_path.with_suffix('.input.flac')
|
||||
output_path = output_cuts_path.parent / input_relative_path.with_suffix('.output.flac')
|
||||
|
||||
# we know that `audio_dir` exists, but we need to create the parent directories
|
||||
input_path.parent.mkdir(exist_ok=True, parents=True)
|
||||
output_path.parent.mkdir(exist_ok=True, parents=True)
|
||||
else:
|
||||
(output_cuts_path.parent / 'audio').mkdir(exist_ok=True, parents=True)
|
||||
input_path = output_cuts_path.parent / 'audio' / f"{i:06}.input.flac"
|
||||
output_path = output_cuts_path.parent / 'audio' / f"{i:06}.output.flac"
|
||||
|
||||
sf.write(input_path, input_audio, sample_rate, format='FLAC', subtype='PCM_24')
|
||||
sf.write(output_path, output_audio, sample_rate, format='FLAC', subtype='PCM_24')
|
||||
|
||||
input_recording = Recording.from_file(input_path)
|
||||
input_recording.sources[0].source = str(input_path.relative_to(output_cuts_path.parent))
|
||||
output_recording = Recording.from_file(output_path)
|
||||
output_recording.sources[0].source = str(output_path.relative_to(output_cuts_path.parent))
|
||||
|
||||
cut = MonoCut(
|
||||
id=input_recording.id, start=0, channel=0, duration=input_recording.duration, recording=input_recording
|
||||
)
|
||||
cut.target_recording = output_recording
|
||||
|
||||
for optional_field_name in (
|
||||
'text',
|
||||
'original_text',
|
||||
'language',
|
||||
):
|
||||
if (
|
||||
hasattr(original_cut, optional_field_name)
|
||||
and getattr(original_cut, optional_field_name) is not None
|
||||
):
|
||||
setattr(cut, optional_field_name, getattr(original_cut, optional_field_name))
|
||||
|
||||
writer.write(cut)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user