run_evaluation.py

# We need to set CUDA_VISIBLE_DEVICES before we import Pytorch so we will read all arguments directly on startup
import logging
import os
from argparse import ArgumentParser
from collections import defaultdict
from pathlib import Path
import pandas as pd
from typing import List

parser = ArgumentParser()
parser.add_argument('--config', default='config_eval.yaml')
parser.add_argument('--gpu_ids', default='0')
args = parser.parse_args()

if 'CUDA_VISIBLE_DEVICES' not in os.environ:  # do not overwrite previously set devices
    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids

import torch
import time
import shutil
import itertools

from evaluation import evaluate_asv, train_asv_eval, evaluate_asr, train_asr_eval, evaluate_gvd
from utils import parse_yaml, find_asv_model_checkpoint, scan_checkpoint, combine_asr_data, split_vctk_into_common_and_diverse


def get_evaluation_steps(params):
    eval_steps = {}
    if 'privacy' in params:
        eval_steps['privacy'] = list(params['privacy'].keys())
    if 'utility' in params:
        eval_steps['utility'] = list(params['utility'].keys())

    if 'eval_steps' in params:
        param_eval_steps = params['eval_steps']

        for eval_part, eval_metrics in param_eval_steps.items():
            if eval_part not in eval_steps:
                raise KeyError(f'Unknown evaluation step {eval_part}, please specify in config')
            for metric in eval_metrics:
                if metric not in eval_steps[eval_part]:
                    raise KeyError(f'Unknown metric {metric}, please specify in config')
        return param_eval_steps
    return eval_steps


def get_eval_trial_datasets(datasets_list):
    eval_pairs = []

    for dataset in datasets_list:
        eval_pairs.extend([(f'{dataset["data"]}_{dataset["set"]}_{enroll}',
                            f'{dataset["data"]}_{dataset["set"]}_{trial}')
                           for enroll, trial in itertools.product(dataset['enrolls'], dataset['trials'])])
    return eval_pairs


def check_vctk_split(data_trials, eval_data_dir, anon_suffix):
    copy_files_for_orig = ['spk2utt', 'text', 'utt2spk', 'trials', 'spk2gender', 'wav.scp', 'utt2dur']
    copy_files_for_anon = ['spk2utt', 'text', 'utt2spk', 'trials']
    separated_data_trials = []

    for enroll, trial in data_trials:
        if 'vctk' in trial and '_all' in trial:
            common_split = trial.replace('all', 'common')  # same sentences for all speakers
            diverse_split = trial.replace('_all', '')  # different sentences for each speaker
            if (not Path(eval_data_dir, common_split).exists() or \
                    not Path(eval_data_dir, diverse_split).exists()):
                split_vctk_into_common_and_diverse(dataset=trial, output_path=eval_data_dir, orig_data_path=eval_data_dir,
                                                   copy_files=copy_files_for_orig, anon=False, anon_suffix=f'_{anon_suffix}',
                                                   out_data_split=Path(eval_data_dir, trial))
            if not Path(eval_data_dir, f'{common_split}_{anon_suffix}').exists() or \
                not Path(eval_data_dir, f'{diverse_split}_{anon_suffix}').exists():
                split_vctk_into_common_and_diverse(dataset=trial, output_path=eval_data_dir, orig_data_path=eval_data_dir,
                                                   copy_files=copy_files_for_anon, anon=True, anon_suffix=f'_{anon_suffix}',
                                                   out_data_split = Path(eval_data_dir, f'{trial}_{anon_suffix}'))
            separated_data_trials.append((enroll, common_split))
            separated_data_trials.append((enroll, diverse_split))
        else:
            separated_data_trials.append((enroll, trial))
    return separated_data_trials


def get_eval_asr_datasets(datasets_list, eval_data_dir, anon_suffix):
    # combines the trial subsets (trial_f, trial_m) of each dataset into one asr dataset
    eval_data = set()

    # in case one bigger dataset is divided into smaller parts (e.g. vctk_all into vctk and vctk_common) we need to
    # collect all parts together to create one asr dataset for the whole instead of for each subpart
    collated_datasets = defaultdict(list)
    for dataset in datasets_list:
        asr_dataset_name = f'{dataset["data"]}_{dataset["set"]}_asr'
        collated_datasets[asr_dataset_name].append(dataset)

    for asr_dataset_name, datasets in collated_datasets.items():
        for asr_dataset in (asr_dataset_name, f'{asr_dataset_name}_{anon_suffix}'): # for orig and anon
            if not Path(eval_data_dir / asr_dataset).exists():
                trial_dirs = []
                for dataset in datasets:
                    for trial in dataset['trials']:
                        trial_dirs.append(Path(eval_data_dir / f'{dataset["data"]}_{dataset["set"]}_{trial}'))
                output_dir = Path(eval_data_dir / asr_dataset)
                combine_asr_data(input_dirs=trial_dirs, output_dir=output_dir)
        eval_data.add(asr_dataset_name)

    return list(eval_data)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s- %(levelname)s - %(message)s')

    params = parse_yaml(Path('configs', args.config))
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    eval_data_dir = params['eval_data_dir']
    anon_suffix = params['anon_data_suffix']
    eval_steps = get_evaluation_steps(params)
    eval_data_trials = get_eval_trial_datasets(params['datasets'])
    eval_data_trials = check_vctk_split(eval_data_trials, eval_data_dir=eval_data_dir, anon_suffix=anon_suffix)
    eval_data_asr = get_eval_asr_datasets(params['datasets'], eval_data_dir=eval_data_dir, anon_suffix=anon_suffix)

    # make sure given paths exist
    assert eval_data_dir.exists(), f'{eval_data_dir} does not exist'

    if 'privacy' in eval_steps:
        if 'asv' in eval_steps['privacy']:
            asv_params = params['privacy']['asv']
            model_dir = find_asv_model_checkpoint(asv_params['model_dir'])
            if 'training' in asv_params:
                asv_train_params = asv_params['training']
                if not model_dir.exists() or asv_train_params.get('retrain', True) is True:
                    start_time = time.time()
                    logging.info('Perform ASV training')
                    train_asv_eval(train_params=asv_train_params, output_dir=asv_params['model_dir'])
                    logging.info("ASV training time: %f min ---" % (float(time.time() - start_time) / 60))
                    model_dir = scan_checkpoint(model_dir, 'CKPT')
                    if asv_params['vec_type'] == 'xvector':
                        shutil.copy('evaluation/privacy/asv/asv_train/hparams/xvector/hyperparams.yaml', model_dir)
                    else:
                        shutil.copy('evaluation/privacy/asv/asv_train/hparams/ecapa/hyperparams.yaml', model_dir)

            if 'evaluation' in asv_params:
                logging.info('Perform ASV evaluation')
                start_time = time.time()
                evaluate_asv(eval_datasets=eval_data_trials, eval_data_dir=eval_data_dir, params=asv_params,
                             device=device,  model_dir=model_dir, anon_data_suffix=anon_suffix)
                logging.info("--- EER computation time: %f min ---" % (float(time.time() - start_time) / 60))

    if 'utility' in eval_steps:
        if 'asr' in eval_steps['utility']:
            asr_params = params['utility']['asr']
            model_name = asr_params['model_name']
            model_dir = asr_params['model_dir']
            libri_dir = asr_params['libri_dir']
            backend = asr_params['backend'].lower()
            asr_model_path = None

            if 'training' in asr_params:
                asr_train_params = asr_params['training']
                asr_model_path = model_dir / 'asr_train_asr_transformer_raw_en_bpe5000'
                if asr_train_params['anon'] and asr_train_params['finetuning']:
                    asr_model_path = model_dir / 'asr_train_asr_transformer_anon_raw_en_bpe5000'

                if not model_dir.exists() or asr_train_params.get('retrain', True) is True:
                    start_time = time.time()
                    logging.info('Perform ASR training')
                    train_asr_eval(params=asr_train_params, libri_dir=libri_dir, model_name=model_name,
                                   model_dir=model_dir, anon_data_suffix=anon_suffix)
                    logging.info("--- ASR training time: %f min ---" % (float(time.time() - start_time) / 60))

            if 'evaluation' in asr_params:
                asr_eval_params = asr_params['evaluation']
                if not asr_model_path:
                    if backend == 'espnet':
                        if asr_eval_params.get('anon_train_model', False):
                            asr_model_path = model_dir / 'asr_train_asr_transformer_anon_raw_en_bpe5000'
                        else:
                            asr_model_path = model_dir / 'asr_train_asr_transformer_raw_en_bpe5000'
                    else:
                        asr_model_path = model_dir

                start_time = time.time()
                logging.info('Perform ASR evaluation')
                evaluate_asr(eval_datasets=eval_data_asr, eval_data_dir=eval_data_dir, params=asr_eval_params,
                             model_path=asr_model_path, anon_data_suffix=anon_suffix, libri_dir=libri_dir,
                             device=device, backend=backend)
                logging.info("--- ASR evaluation time: %f min ---" % (float(time.time() - start_time) / 60))

        if 'gvd' in eval_steps['utility']:
            gvd_params = params['utility']['gvd']
            start_time = time.time()
            logging.info('Perform GVD evaluation')
            evaluate_gvd(eval_datasets=eval_data_trials, eval_data_dir=eval_data_dir, params=gvd_params,
                         device=device, anon_data_suffix=anon_suffix)
            logging.info("--- GVD  computation time: %f min ---" % (float(time.time() - start_time) / 60))