|
1 | 1 | #!/usr/bin/python |
| 2 | + |
| 3 | +""" |
| 4 | +Recipe to train CONV-TASNET model on the WSJ0 dataset |
| 5 | +
|
| 6 | +Author: |
| 7 | + * Cem Subakan 2020 |
| 8 | +""" |
| 9 | + |
2 | 10 | import os |
3 | 11 | import speechbrain as sb |
4 | 12 | from speechbrain.utils.train_logger import summarize_average |
5 | 13 | import torch |
6 | 14 | from speechbrain.utils.checkpoints import ckpt_recency |
7 | 15 | from speechbrain.nnet.losses import get_si_snr_with_pitwrapper |
| 16 | +from prepare_data import create_wsj_csv, get_wsj_files |
8 | 17 |
|
9 | 18 | import torch.nn.functional as F |
10 | | -import csv |
11 | 19 |
|
12 | 20 | experiment_dir = os.path.dirname(os.path.realpath(__file__)) |
13 | 21 | params_file = os.path.join(experiment_dir, "params.yaml") |
|
16 | 24 | params = sb.yaml.load_extended_yaml(fin) |
17 | 25 |
|
18 | 26 | # this points to the folder which holds the wsj0-mix dataset folder |
19 | | -datapath = params.datapath |
| 27 | +wsj0root = params.wsj0path |
| 28 | +data_save_dir = params.datapath |
20 | 29 |
|
| 30 | +if os.path.exists(data_save_dir): |
| 31 | + get_wsj_files(wsj0root, data_save_dir) |
21 | 32 |
|
22 | 33 | # load or create the csv files for the data |
23 | 34 | if not ( |
24 | 35 | os.path.exists("wsj_tr.csv") |
25 | 36 | and os.path.exists("wsj_cv.csv") |
26 | 37 | and os.path.exists("wsj_tt.csv") |
27 | 38 | ): |
28 | | - for set_type in ["tr", "cv", "tt"]: |
29 | | - mix_path = ( |
30 | | - datapath + "wsj0-mix/2speakers/wav8k/min/" + set_type + "/mix/" |
31 | | - ) |
32 | | - s1_path = datapath + "wsj0-mix/2speakers/wav8k/min/" + set_type + "/s1/" |
33 | | - s2_path = datapath + "wsj0-mix/2speakers/wav8k/min/" + set_type + "/s2/" |
34 | | - |
35 | | - files = os.listdir(mix_path) |
36 | | - |
37 | | - mix_fl_paths = [mix_path + fl for fl in files] |
38 | | - s1_fl_paths = [s1_path + fl for fl in files] |
39 | | - s2_fl_paths = [s2_path + fl for fl in files] |
40 | | - |
41 | | - csv_columns = [ |
42 | | - "ID", |
43 | | - "duration", |
44 | | - "mix_wav", |
45 | | - "mix_wav_format", |
46 | | - "mix_wav_opts", |
47 | | - "s1_wav", |
48 | | - "s1_wav_format", |
49 | | - "s1_wav_opts", |
50 | | - "s2_wav", |
51 | | - "s2_wav_format", |
52 | | - "s2_wav_opts", |
53 | | - ] |
54 | | - |
55 | | - with open("wsj_" + set_type + ".csv", "w") as csvfile: |
56 | | - writer = csv.DictWriter(csvfile, fieldnames=csv_columns) |
57 | | - writer.writeheader() |
58 | | - for i, (mix_path, s1_path, s2_path) in enumerate( |
59 | | - zip(mix_fl_paths, s1_fl_paths, s2_fl_paths) |
60 | | - ): |
61 | | - |
62 | | - row = { |
63 | | - "ID": i, |
64 | | - "duration": 1.0, |
65 | | - "mix_wav": mix_path, |
66 | | - "mix_wav_format": "wav", |
67 | | - "mix_wav_opts": None, |
68 | | - "s1_wav": s1_path, |
69 | | - "s1_wav_format": "wav", |
70 | | - "s1_wav_opts": None, |
71 | | - "s2_wav": s2_path, |
72 | | - "s2_wav_format": "wav", |
73 | | - "s2_wav_opts": None, |
74 | | - } |
75 | | - writer.writerow(row) |
76 | | - |
| 39 | + create_wsj_csv(data_save_dir) |
77 | 40 |
|
78 | 41 | tr_csv = os.path.realpath(os.path.join(experiment_dir, "wsj_tr.csv")) |
79 | 42 | cv_csv = os.path.realpath(os.path.join(experiment_dir, "wsj_cv.csv")) |
|
85 | 48 | ) |
86 | 49 | # print(params) # if needed this line can be uncommented for logging |
87 | 50 |
|
88 | | - |
89 | 51 | if params.use_tensorboard: |
90 | 52 | from speechbrain.utils.train_logger import TensorboardLogger |
91 | 53 |
|
92 | 54 | train_logger = TensorboardLogger(params.tensorboard_logs) |
93 | 55 |
|
94 | | - |
95 | 56 | device = "cuda" if torch.cuda.is_available() else "cpu" |
96 | 57 |
|
97 | 58 |
|
@@ -200,7 +161,6 @@ def on_epoch_end(self, epoch, train_stats, valid_stats): |
200 | 161 |
|
201 | 162 | params.checkpointer.recover_if_possible(lambda c: -c.meta["av_loss"]) |
202 | 163 |
|
203 | | -# with torch.autograd.detect_anomaly(): |
204 | 164 | ctn.fit( |
205 | 165 | range(params.N_epochs), |
206 | 166 | train_set=train_loader, |
|
0 commit comments