Skip to content

Commit f0fe561

Browse files
mravanellimravanelli
authored
Merge pull request speechbrain#482 from speechbrain/MetricGAN+
MetricGAN+
2 parents 7bdebdc + 861c265 commit f0fe561

8 files changed

Lines changed: 899 additions & 4 deletions

File tree

‎recipes/Voicebank/enhance/MetricGAN/README.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# MetricGAN Recipe for Enhancement
2+
3+
This recipe implements MetricGAN recipe for enhancement as described in the paper
4+
[MetricGAN: Generative Adversarial Networks based Black-box Metric Scores Optimization for Speech Enhancement](https://arxiv.org/abs/1905.04874)
5+
6+
Use the `download_vctk` function in `voicebank_prepare.py` to download the dataset
7+
and resample it to 16000 Hz. To run an experiment, execute the following command in
8+
the current folder:
9+
10+
```bash
11+
python train.py hparams/train.yaml --data_folder /path/to/data_folder
12+
```
13+
14+
## Results
15+
16+
Experiment Date | PESQ | STOI
17+
-|-|-
18+
2021-03-06 | 3.08 | 93.0
19+
20+
## Citation
21+
22+
If you find the code useful in your research, please cite:
23+
24+
@inproceedings{fu2019metricGAN,
25+
title = {MetricGAN: Generative Adversarial Networks based Black-box Metric Scores Optimization for Speech Enhancement},
26+
author = {Fu, Szu-Wei and Liao, Chien-Feng and Tsao, Yu and Lin, Shou-De},
27+
booktitle = {International Conference on Machine Learning (ICML)},
28+
year = {2019}
29+
}

‎recipes/Voicebank/enhance/MetricGAN/hparams/train.yaml‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
# #################################
2+
# Basic training parameters
3+
# To train a different model, change "!include:" statement to new model file
4+
# To compute loss in the time domain, switch "waveform_target" to True
5+
# Authors:
6+
# * Szu-Wei Fu 2021
7+
# * Peter Plantinga 2020, 2021
8+
# #################################
9+
10+
# Seed needs to be set at top of yaml, before objects with parameters are made
11+
seed: 7230
12+
__set_seed: !!python/object/apply:torch.manual_seed [!ref <seed>]
13+
14+
data_folder: !PLACEHOLDER # e.g, /data/member1/user_jasonfu/noisy-vctk-16k
15+
train_clean_folder: !ref <data_folder>/clean_trainset_28spk_wav_16k/
16+
train_noisy_folder: !ref <data_folder>/noisy_trainset_28spk_wav_16k/
17+
valid_clean_folder: !ref <data_folder>/clean_testset_wav_16k/
18+
test_clean_folder: !ref <data_folder>/clean_testset_wav_16k/
19+
20+
MetricGAN_folder: !ref <output_folder>/enhanced_wavs
21+
output_folder: !ref ./results/MetricGAN/<seed>
22+
save_folder: !ref <output_folder>/save
23+
train_log: !ref <output_folder>/train_log.txt
24+
enhanced_folder: !ref <output_folder>/enhanced_wavs
25+
26+
# Basic parameters
27+
use_tensorboard: False
28+
tensorboard_logs: !ref <output_folder>/logs/
29+
30+
# FFT paremeters
31+
Sample_rate: 16000
32+
Win_length: 32
33+
Hop_length: 16
34+
N_fft: 512
35+
window_fn: !name:torch.hamming_window
36+
37+
# Data files
38+
train_annotation: !ref <save_folder>/train.json
39+
valid_annotation: !ref <save_folder>/valid.json
40+
test_annotation: !ref <save_folder>/test.json
41+
skip_prep: False
42+
43+
# The target metrics that you want to optimize.
44+
# Right now we only support 'pesq' and 'stoi'.
45+
# (Of course, it can be any arbitary metric.)
46+
target_metric: pesq
47+
# Training Parameters
48+
number_of_epochs: 600
49+
number_of_samples: 100
50+
min_mask: 0
51+
train_N_batch: 1
52+
valid_N_batch: 30
53+
history_portion: 0.2
54+
G_lr: 0.0005
55+
D_lr: 0.0005
56+
mse_weight: 0
57+
58+
dataloader_options:
59+
batch_size: !ref <train_N_batch>
60+
valid_dataloader_options:
61+
batch_size: !ref <valid_N_batch>
62+
63+
# Change this import to use a different model
64+
models: !include:../models/MetricGAN.yaml
65+
N_fft: !ref <N_fft>
66+
67+
epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
68+
limit: !ref <number_of_epochs>
69+
70+
modules:
71+
generator: !ref <models[generator]>
72+
discriminator: !ref <models[discriminator]>
73+
74+
g_opt_class: !name:torch.optim.Adam
75+
lr: !ref <G_lr>
76+
d_opt_class: !name:torch.optim.Adam
77+
lr: !ref <D_lr>
78+
79+
checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
80+
checkpoints_dir: !ref <save_folder>
81+
recoverables:
82+
model: !ref <models[generator]>
83+
counter: !ref <epoch_counter>
84+
85+
compute_cost: !name:speechbrain.nnet.losses.mse_loss
86+
87+
compute_STFT: !new:speechbrain.processing.features.STFT
88+
sample_rate: !ref <Sample_rate>
89+
win_length: !ref <Win_length>
90+
hop_length: !ref <Hop_length>
91+
n_fft: !ref <N_fft>
92+
window_fn: !ref <window_fn>
93+
94+
compute_ISTFT: !new:speechbrain.processing.features.ISTFT
95+
sample_rate: !ref <Sample_rate>
96+
win_length: !ref <Win_length>
97+
hop_length: !ref <Hop_length>
98+
window_fn: !ref <window_fn>
99+
100+
resynth: !name:speechbrain.processing.signal_processing.resynthesize
101+
stft: !ref <compute_STFT>
102+
istft: !ref <compute_ISTFT>
103+
normalize_wavs: False
104+
105+
train_logger: !new:speechbrain.utils.train_logger.FileTrainLogger
106+
save_file: !ref <train_log>

‎recipes/Voicebank/enhance/MetricGAN/models/MetricGAN.yaml‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# ################################
2+
# Model: MetricGAN
3+
# Authors: Szu-Wei Fu 2021
4+
# ################################
5+
6+
# Neural parameters
7+
kernel_size: (5,5)
8+
base_channels: 15
9+
10+
N_fft: !PLACEHOLDER
11+
12+
generator: !new:speechbrain.lobes.models.MetricGAN.EnhancementGenerator
13+
14+
discriminator: !new:speechbrain.lobes.models.MetricGAN.MetricDiscriminator
15+
kernel_size: !ref <kernel_size>
16+
base_channels: !ref <base_channels>

0 commit comments

Comments
 (0)