NUM_WORKERS = 5 train_loader = DataLoader(timit_train, shuffle=True, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, drop_last=True) val_loader = DataLoader(timit_val, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, drop_last=True) # some random parameters, does it look sensible? LR = 1e-3 REDUCE_LR_PATIENCE = 5 EARLY_STOP_PATIENCE = 20 MAX_EPOCHS = 300 # the model here should be constructed in the script accordingly to the passed config (including the model type) # most of the models accept `sample_rate` parameter for encoders, which is important (default is 16000, override) #model = DCUNet("DCUNet-20", fix_length_mode="trim", sample_rate=SAMPLE_RATE) model = ConvTasNet(n_src=1) checkpoint = ModelCheckpoint( filename='{epoch:02d}-{val_loss:.2f}', monitor="val_loss", mode="min", save_top_k=5, verbose=True ) optimizer = optim.Adam(model.parameters(), lr=LR) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=REDUCE_LR_PATIENCE) early_stopping = EarlyStopping(monitor='val_loss', patience=EARLY_STOP_PATIENCE) # Probably we also need to subclass `System`, in order to log the target metrics on the validation set (PESQ/STOI) system = System(model, optimizer, sisdr_loss_wrapper, train_loader, val_loader, scheduler)
def main(conf): model_path = os.path.join(conf['exp_dir'], 'best_model.pth') model = ConvTasNet.from_pretrained(model_path) # Handle device placement if conf['use_gpu']: model.cuda() model_device = next(model.parameterss()).device test_set = LibriMix(csv_dir=conf['test_dir'], task=conf['task'], sample_rate=conf['sample_rate'], n_src=conf['train_conf']['data']['n_src'], segment=None) # Uses all segment length # Used to reorder sources only loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from='pw_mtx') # Randomly choose the indexes of sentences to save. eval_save_dir = os.path.join(conf['exp_dir'], conf['out_dir']) ex_save_dir = os.path.join(eval_save_dir, 'examples/') if conf['n_save_ex'] == -1: conf['n_save_ex'] = len(test_set) save_idx = random.sample(range(len(test_set)), conf['n_save_ex']) series_list = [] torch.no_grad().__enter__() for idx in tqdm(range(len(test_set))): # Forward the network on the mixture. mix, sources = tensors_to_device(test_set[idx], device=model_device) est_sources = model(mix.unsqueeze(0)) loss, reordered_sources = loss_func(est_sources, sources[None], return_est=True) mix_np = mix.cpu().data.numpy() sources_np = sources.squeeze().cpu().data.numpy() est_sources_np = reordered_sources.squeeze().cpu().data.numpy() # For each utterance, we get a dictionary with the mixture path, # the input and output metrics utt_metrics = get_metrics(mix_np, sources_np, est_sources_np, sample_rate=conf['sample_rate']) utt_metrics['mix_path'] = test_set.mixture_path series_list.append(pd.Series(utt_metrics)) # Save some examples in a folder. Wav files and metrics as text. if idx in save_idx: local_save_dir = os.path.join(ex_save_dir, 'ex_{}/'.format(idx)) os.makedirs(local_save_dir, exist_ok=True) sf.write(local_save_dir + "mixture.wav", mix_np, conf['sample_rate']) # Loop over the sources and estimates for src_idx, src in enumerate(sources_np): sf.write(local_save_dir + "s{}.wav".format(src_idx), src, conf['sample_rate']) for src_idx, est_src in enumerate(est_sources_np): sf.write(local_save_dir + "s{}_estimate.wav".format(src_idx), est_src, conf['sample_rate']) # Write local metrics to the example folder. with open(local_save_dir + 'metrics.json', 'w') as f: json.dump(utt_metrics, f, indent=0) # Save all metrics to the experiment folder. all_metrics_df = pd.DataFrame(series_list) all_metrics_df.to_csv(os.path.join(eval_save_dir, 'all_metrics.csv')) # Print and save summary metrics final_results = {} for metric_name in compute_metrics: input_metric_name = 'input_' + metric_name ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name] final_results[metric_name] = all_metrics_df[metric_name].mean() final_results[metric_name + '_imp'] = ldf.mean() print('Overall metrics :') pprint(final_results) with open(os.path.join(eval_save_dir, 'final_metrics.json'), 'w') as f: json.dump(final_results, f, indent=0) model_dict = torch.load(model_path, map_location='cpu') publishable = save_publishable(os.path.join(conf['exp_dir'], 'publish_dir'), model_dict, metrics=final_results, train_conf=train_conf)
total_df.sort_values(['SI-SDR', 'PESQ', 'STOI'], inplace=True) total_df = total_df.round({'SI-SDR': 3, 'PESQ': 3, 'STOI': 3}) print(total_df) return total_df models = { 'input': None, 'baseline': RegressionFCNN.from_pretrained('models/baseline_model_v1.pt'), 'vae': VAE.from_pretrained('/jmain01/home/JAD007/txk02/aaa18-txk02/workspace/models/VAE.pt'), 'auto_encoder': VAE.from_pretrained('/jmain01/home/JAD007/txk02/aaa18-txk02/workspace/models/AutoEncoder.pt'), 'waveunet_v1': WaveUNet.from_pretrained('models/waveunet_model_adapt.pt'), 'dcunet_20': DCUNet.from_pretrained('models/dcunet_20_random_v2.pt'), 'dccrn': DCCRNet.from_pretrained('models/dccrn_random_v1.pt'), 'smolnet': SMoLnet.from_pretrained('models/SMoLnet.pt'), 'dprnn': DPRNNTasNet.from_pretrained('models/dprnn_model.pt'), 'conv_tasnet': ConvTasNet.from_pretrained('models/convtasnet_model.pt'), 'dptnet': DPTNet.from_pretrained('models/dptnet_model.pt'), 'demucs': Demucs.from_pretrained('models/Demucs.pt'), } def eval_all_and_plot(models, test_set, directory, plot_name): results_dfs = {} for model_name, model in models.items(): print(f'Evaluating {model_labels[model_name]}') csv_path = f'/jmain01/home/JAD007/txk02/aaa18-txk02/DRONE_project/asteroid/notebooks/{directory}/{model_name}.csv' if os.path.isfile(csv_path): print('Results already available') df = pd.read_csv(csv_path) else:
def main(conf): model_path = os.path.join(conf["exp_dir"], "best_model.pth") model = ConvTasNet.from_pretrained(model_path) model = LambdaOverlapAdd( nnet=model, # function to apply to each segment. n_src=2, # number of sources in the output of nnet window_size=64000, # Size of segmenting window hop_size=None, # segmentation hop size window="hanning", # Type of the window (see scipy.signal.get_window reorder_chunks=False, # Whether to reorder each consecutive segment. enable_grad= False, # Set gradient calculation on of off (see torch.set_grad_enabled) ) # Handle device placement if conf["use_gpu"]: model.cuda() model_device = next(model.parameters()).device # Evaluation is mode using 'remix' mixture dataset_kwargs = { "root_path": Path(conf["train_conf"]["data"]["root_path"]), "task": conf["train_conf"]["data"]["task"], "sample_rate": conf["train_conf"]["data"]["sample_rate"], "num_workers": conf["train_conf"]["training"]["num_workers"], "mixture": "remix", } test_set = DAMPVSEPDataset(split="test", **dataset_kwargs) # Randomly choose the indexes of sentences to save. eval_save_dir = os.path.join(conf["exp_dir"], conf["out_dir"]) ex_save_dir = os.path.join(eval_save_dir, "examples/") if conf["n_save_ex"] == -1: conf["n_save_ex"] = len(test_set) save_idx = random.sample(range(len(test_set)), conf["n_save_ex"]) series_list = [] torch.no_grad().__enter__() for idx in tqdm(range(len(test_set))): # Forward the network on the mixture. mix, sources = test_set[idx] mix = mix.to(model_device) est_sources = model.forward(mix.unsqueeze(0).unsqueeze(1)) mix_np = mix.squeeze(0).cpu().data.numpy() sources_np = sources.cpu().data.numpy() est_sources_np = est_sources.squeeze(0).cpu().data.numpy() # For each utterance, we get a dictionary with the mixture path, # the input and output metrics utt_metrics = get_metrics( mix_np, sources_np, est_sources_np, sample_rate=conf["sample_rate"], metrics_list=compute_metrics, average=False, ) utt_metrics = split_metric_dict(utt_metrics) utt_metrics["mix_path"] = test_set.mixture_path series_list.append(pd.Series(utt_metrics)) # Save some examples in a folder. Wav files and metrics as text. if idx in save_idx: local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx)) os.makedirs(local_save_dir, exist_ok=True) sf.write(local_save_dir + "mixture.wav", mix_np / max(abs(mix_np)), conf["sample_rate"]) # Loop over the sources and estimates for src_idx, src in enumerate(sources_np): sf.write(local_save_dir + "s{}.wav".format(src_idx), src, conf["sample_rate"]) for src_idx, est_src in enumerate(est_sources_np): est_src *= np.max(np.abs(mix_np)) / np.max(np.abs(est_src)) sf.write( local_save_dir + "s{}_estimate.wav".format(src_idx), est_src, conf["sample_rate"], ) # Write local metrics to the example folder. with open(local_save_dir + "metrics.json", "w") as f: json.dump(utt_metrics, f, indent=0) # Save all metrics to the experiment folder. all_metrics_df = pd.DataFrame(series_list) all_metrics_df.to_csv(os.path.join(eval_save_dir, "all_metrics.csv")) # Print and save summary metrics final_results = {} for metric_name in compute_metrics: for s in ["", "_s0", "_s1"]: input_metric_name = "input_" + f"{metric_name}{s}" ldf = all_metrics_df[f"{metric_name}{s}"] - all_metrics_df[ input_metric_name] final_results[f"{metric_name}{s}"] = all_metrics_df[ f"{metric_name}{s}"].mean() final_results[f"{metric_name}{s}" + "_imp"] = ldf.mean() print("Overall metrics :") pprint(final_results) with open(os.path.join(eval_save_dir, "final_metrics.json"), "w") as f: json.dump(final_results, f, indent=0) model_dict = torch.load(model_path, map_location="cpu") os.makedirs(os.path.join(conf["exp_dir"], "publish_dir"), exist_ok=True) publishable = save_publishable( os.path.join(conf["exp_dir"], "publish_dir"), model_dict, metrics=final_results, train_conf=train_conf, )
def main(conf): compute_metrics = update_compute_metrics(conf["compute_wer"], COMPUTE_METRICS) anno_df = pd.read_csv(Path(conf["test_dir"]).parent.parent.parent / "test_annotations.csv") wer_tracker = ( MockWERTracker() if not conf["compute_wer"] else WERTracker(ASR_MODEL_PATH, anno_df) ) model_path = os.path.join(conf["exp_dir"], "best_model.pth") model = ConvTasNet.from_pretrained(model_path) # Handle device placement if conf["use_gpu"]: model.cuda() model_device = next(model.parameters()).device test_set = LibriMix( csv_dir=conf["test_dir"], task=conf["task"], sample_rate=conf["sample_rate"], n_src=conf["train_conf"]["data"]["n_src"], segment=None, return_id=True, ) # Uses all segment length # Used to reorder sources only loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx") # Randomly choose the indexes of sentences to save. eval_save_dir = os.path.join(conf["exp_dir"], conf["out_dir"]) ex_save_dir = os.path.join(eval_save_dir, "examples/") if conf["n_save_ex"] == -1: conf["n_save_ex"] = len(test_set) save_idx = random.sample(range(len(test_set)), conf["n_save_ex"]) series_list = [] torch.no_grad().__enter__() for idx in tqdm(range(len(test_set))): # Forward the network on the mixture. mix, sources, ids = test_set[idx] mix, sources = tensors_to_device([mix, sources], device=model_device) est_sources = model(mix.unsqueeze(0)) loss, reordered_sources = loss_func(est_sources, sources[None], return_est=True) mix_np = mix.cpu().data.numpy() sources_np = sources.cpu().data.numpy() est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy() # For each utterance, we get a dictionary with the mixture path, # the input and output metrics utt_metrics = get_metrics( mix_np, sources_np, est_sources_np, sample_rate=conf["sample_rate"], metrics_list=COMPUTE_METRICS, ) utt_metrics["mix_path"] = test_set.mixture_path utt_metrics.update( **wer_tracker( mix=mix_np, clean=sources_np, estimate=est_sources_np, wav_id=ids, sample_rate=conf["sample_rate"], ) ) series_list.append(pd.Series(utt_metrics)) # Save some examples in a folder. Wav files and metrics as text. if idx in save_idx: local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx)) os.makedirs(local_save_dir, exist_ok=True) sf.write(local_save_dir + "mixture.wav", mix_np, conf["sample_rate"]) # Loop over the sources and estimates for src_idx, src in enumerate(sources_np): sf.write(local_save_dir + "s{}.wav".format(src_idx), src, conf["sample_rate"]) for src_idx, est_src in enumerate(est_sources_np): est_src *= np.max(np.abs(mix_np)) / np.max(np.abs(est_src)) sf.write( local_save_dir + "s{}_estimate.wav".format(src_idx), est_src, conf["sample_rate"], ) # Write local metrics to the example folder. with open(local_save_dir + "metrics.json", "w") as f: json.dump(utt_metrics, f, indent=0) # Save all metrics to the experiment folder. all_metrics_df = pd.DataFrame(series_list) all_metrics_df.to_csv(os.path.join(eval_save_dir, "all_metrics.csv")) # Print and save summary metrics final_results = {} for metric_name in compute_metrics: input_metric_name = "input_" + metric_name ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name] final_results[metric_name] = all_metrics_df[metric_name].mean() final_results[metric_name + "_imp"] = ldf.mean() print("Overall metrics :") pprint(final_results) if conf["compute_wer"]: print("\nWER report") wer_card = wer_tracker.final_report_as_markdown() print(wer_card) # Save the report with open(os.path.join(eval_save_dir, "final_wer.md"), "w") as f: f.write(wer_card) with open(os.path.join(eval_save_dir, "final_metrics.json"), "w") as f: json.dump(final_results, f, indent=0) model_dict = torch.load(model_path, map_location="cpu") os.makedirs(os.path.join(conf["exp_dir"], "publish_dir"), exist_ok=True) publishable = save_publishable( os.path.join(conf["exp_dir"], "publish_dir"), model_dict, metrics=final_results, train_conf=train_conf, )