Evaluating Models on different number of sources

[anontiger]

I have managed to train my model on a dataset with three sources. I was wondering, is it possible to evaluate how this will perform on two sources? As I trained it differently from run.sh, I had to use a modified eval.py as well. Here are my train.py and eval.py:

train.py
from torch import optim
from pytorch_lightning import Trainer
from asteroid.losses import pairwise_neg_sisdr, PITLossWrapper
from asteroid.data import LibriMix
from asteroid.engine import System
from asteroid.models import DPRNNTasNet
from customdata import customdata
import torch
import torchvision

#train_loader, val_loader = DPRNNTasNet.loaders_from_mini(task="sep_clean", batch_size=16)
train_loader, val_loader,test_loader = customdata.loaders_from_mini(task="sep_clean", n_src=3, batch_size=16)
model = DPRNNTasNet(n_src=3,sampling_rate=22050)

loss = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx")

optimizer = optim.Adam(model.parameters(), lr=1e-3)

system = System(model, optimizer, loss, train_loader, val_loader)

trainer = Trainer(max_epochs=200, gpus=1)

trainer.fit(system)

model.separate("trclean1.wav", force_overwrite=True,resample=True)

torch.save(model, "model1.pth")
torch.save(model, 'model1.pt')

torch.save(model.state_dict(), "state_dict.pth")

eval.py
import os
import random
import soundfile as sf
import torch
import yaml
import json
import argparse
import pandas as pd
from tqdm import tqdm
from pprint import pprint
from pathlib import Path

from asteroid.metrics import get_metrics
from asteroid.data.librimix_dataset import LibriMix
from asteroid.losses import PITLossWrapper, pairwise_neg_sisdr
from asteroid.models import SuDORMRFImprovedNet
from asteroid.models import save_publishable
from asteroid.utils import tensors_to_device
from asteroid.dsp.normalization import normalize_estimates
from asteroid.metrics import WERTracker, MockWERTracker
from customdata import customdata
from torch import optim
from pytorch_lightning import Trainer
from asteroid.losses import pairwise_neg_sisdr, PITLossWrapper
from asteroid.data import LibriMix
from asteroid.engine import System
from asteroid.models import DPRNNTasNet
from asteroid.models import SuDORMRFImprovedNet
from customdata import customdata
import torch
import torchvision

parser = argparse.ArgumentParser()
parser.add_argument(
"--test_dir", type=str, required=True, help="Test directory including the csv files"
)
parser.add_argument(
"--task",
type=str,
required=True,
help="One of enh_single, enh_both, " "sep_clean or sep_noisy",
)
parser.add_argument(
"--out_dir",
type=str,
required=True,
help="Directory in exp_dir where the eval results" " will be stored",
)
parser.add_argument(
"--use_gpu", type=int, default=0, help="Whether to use the GPU for model execution"
)
parser.add_argument("--exp_dir", default="exp/tmp", help="Experiment root")
parser.add_argument(
"--n_save_ex", type=int, default=10, help="Number of audio examples to save, -1 means all"
)
parser.add_argument(
"--compute_wer", type=int, default=0, help="Compute WER using ESPNet's pretrained model"
)

COMPUTE_METRICS = ["si_sdr", "sdr", "sir", "sar", "stoi"]
ASR_MODEL_PATH = (
"Shinji Watanabe/librispeech_asr_train_asr_transformer_e18_raw_bpe_sp_valid.acc.best"
)

def update_compute_metrics(compute_wer, metric_list):
if not compute_wer:
return metric_list
try:
from espnet2.bin.asr_inference import Speech2Text
from espnet_model_zoo.downloader import ModelDownloader
except ModuleNotFoundError:
import warnings

    warnings.warn("Couldn't find espnet installation. Continuing without.")
    return metric_list
return metric_list + ["wer"]

def main(conf):
train_loader, val_loader,test_loader = customdata.loaders_from_mini(task="sep_clean", n_src=2, batch_size=4)
model = DPRNNTasNet(n_src=3)
model.load_state_dict(torch.load("state_dict.pth"))

compute_metrics = update_compute_metrics(conf["compute_wer"], COMPUTE_METRICS)
anno_df = pd.read_csv("/mnt/c/asteroid/egs/librimix/SuDORMRFImprovedNet/MiniLibriMix/metadata/mixture_test_mix_clean.csv")
wer_tracker = (
    MockWERTracker() if not conf["compute_wer"] else WERTracker(ASR_MODEL_PATH, anno_df)
)
model_path = "best_model.pth"
# Handle device placement
if conf["use_gpu"]:
    model.cuda()
model_device = next(model.parameters()).device
test_set = customdata(
    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.
    print(test_set[idx])
    print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    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
    est_sources_np_normalized = normalize_estimates(est_sources_np, mix_np)
    utt_metrics.update(
        **wer_tracker(
            mix=mix_np,
            clean=sources_np,
            estimate=est_sources_np_normalized,
            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_normalized):
            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")
#print(model_dict)
#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,
#)

if name == "main":
args = parser.parse_args()
arg_dic = dict(vars(args))
# Load training config
conf_path = "./local/conf.yml"#os.path.join(args.exp_dir, "conf.yml")
with open(conf_path) as f:
train_conf = yaml.safe_load(f)
arg_dic["sample_rate"] = train_conf["data"]["sample_rate"]
arg_dic["train_conf"] = train_conf

if args.task != arg_dic["train_conf"]["data"]["task"]:
    print(
        "Warning : the task used to test is different than "
        "the one from training, be sure this is what you want."
    )
main(arg_dic)

When I have the n_src match up, they train and evaluate just fine, but when I try to evaluate on two sources with a model trained for three sources, I get the following error:
DPRNNTasNetbase.py", line 3621, in get_loc
return self._engine.get_loc(casted_key)
File "pandas/_libs/index.pyx", line 136, in pandas._libs.index.IndexEngine.get_loc
File "pandas/_libs/index.pyx", line 163, in pandas._libs.index.IndexEngine.get_loc
File "pandas/_libs/hashtable_class_helper.pxi", line 5198, in pandas._libs.hashtable.PyObjectHashTable.get_item
File "pandas/_libs/hashtable_class_helper.pxi", line 5206, in pandas._libs.hashtable.PyObjectHashTable.get_item
KeyError: 'source_3_path'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/mnt/c/asteroid/egs/librimix/DPRNNTasNet/eval2.py", line 220, in
main(arg_dic)
File "/mnt/c/asteroid/egs/librimix/DPRNNTasNet/eval2.py", line 119, in main
print(test_set[idx])
File "/mnt/c/asteroid/egs/librimix/DPRNNTasNet/customdata.py", line 105, in getitem
source_path = row[f"source_{i + 1}_path"]
File "/home/user/cd/envs/asteroid2/lib/python3.10/site-packages/pandas/core/series.py", line 958, in getitem
return self._get_value(key)
File "/home/user/cd/envs/asteroid2/lib/python3.10/site-packages/pandas/core/series.py", line 1069, in _get_value
loc = self.index.get_loc(label)
File "/home/user/cd/envs/asteroid2/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 3623, in get_loc
raise KeyError(key) from err
KeyError: 'source_3_path'

I kind of understand that the model expects a third source, but is it impossible to evaluate how it would perform when tested on mixtures with two sources?

[mpariente]

I think you should try to read stack trace, debug your code, etc.. Try to understand the processing flow, which object expects which other object etc.. and why it fails. You'll progress much faster than asking the answer as soon as the problem emerges.

Plus, the error is raised in your custom code, not in Asteroid's.