#!/usr/bin/python #-*- coding: utf-8 -*- import sys, time, os, argparse, socket import yaml import numpy import pdb import torch import glob import zipfile import warnings import datetime from tuneThreshold import * from SpeakerNet import * from DatasetLoader import * import torch.distributed as dist import torch.multiprocessing as mp from scipy.stats import norm from sklearn.mixture import GaussianMixture ## ===== ===== ===== ===== ===== ===== ===== ===== ## Parse arguments ## ===== ===== ===== ===== ===== ===== ===== ===== # os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3' parser = argparse.ArgumentParser(description = "SpeakerNet"); parser.add_argument('--config', type=str, default=None, help='Config YAML file'); ## Data loader parser.add_argument('--max_frames', type=int, default=200, help='Input length to the network for training'); parser.add_argument('--eval_frames', type=int, default=300, help='Input length to the network for testing; 0 uses the whole files'); parser.add_argument('--batch_size', type=int, default=400, help='Batch size, number of speakers per batch'); parser.add_argument('--max_seg_per_spk', type=int, default=500, help='Maximum number of utterances per speaker per epoch'); parser.add_argument('--nDataLoaderThread', type=int, default=10, help='Number of loader threads'); parser.add_argument('--augment', type=bool, default=True, help='Augment input') parser.add_argument('--seed', type=int, default=20211202, help='Seed for the random number generator'); ## Training details parser.add_argument('--test_interval', type=int, default=1, help='Test and save every [test_interval] epochs'); parser.add_argument('--max_epoch', type=int, default=50, help='Maximum number of epochs'); parser.add_argument('--trainfunc', type=str, default="aamsoftmax", help='Loss function'); ## Optimizer parser.add_argument('--optimizer', type=str, default="adamw", help='sgd or adam'); parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler'); parser.add_argument('--lr', type=float, default=0.001, help='Learning rate'); parser.add_argument("--lr_decay", type=float, default=0.9, help='Learning rate decay every [test_interval] epochs'); ## Pre-trained Transformer Model parser.add_argument('--pretrained_model_path', type=str, default="None", help='Absolute path to the pre-trained model'); parser.add_argument('--weight_finetuning_reg', type=float, default=0.001, help='L2 regularization towards the initial pre-trained model'); parser.add_argument('--LLRD_factor', type=float, default=1.0, help='Layer-wise Learning Rate Decay (LLRD) factor'); parser.add_argument('--LR_Transformer', type=float, default=2e-5, help='Learning rate of pre-trained model'); parser.add_argument('--LR_MHFA', type=float, default=5e-3, help='Learning rate of back-end attentive pooling model'); ## Loss functions parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions'); parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions'); parser.add_argument('--margin', type=float, default=0.2, help='Loss margin, only for some loss functions'); parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions'); parser.add_argument('--nPerSpeaker', type=int, default=1, help='Number of utterances per speaker per batch, only for metric learning based losses'); parser.add_argument('--nClasses', type=int, default=5994, help='Number of speakers in the softmax layer, only for softmax-based losses'); ## Evaluation parameters parser.add_argument('--dcf_p_target', type=float, default=0.05, help='A priori probability of the specified target speaker'); parser.add_argument('--dcf_c_miss', type=float, default=1, help='Cost of a missed detection'); parser.add_argument('--dcf_c_fa', type=float, default=1, help='Cost of a spurious detection'); ## Load and save parser.add_argument('--initial_model', type=str, default="", help='Initial model weights'); parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs'); ## Training and test data parser.add_argument('--train_list', type=str, default="data/train_list.txt", help='Train list'); parser.add_argument('--test_list', type=str, default="data/test_list.txt", help='Evaluation list'); parser.add_argument('--train_path', type=str, default="data/voxceleb2", help='Absolute path to the train set'); parser.add_argument('--test_path', type=str, default="data/voxceleb1", help='Absolute path to the test set'); parser.add_argument('--musan_path', type=str, default="data/musan_split", help='Absolute path to the test set'); parser.add_argument('--rir_path', type=str, default="data/simulated_rirs", help='Absolute path to the test set'); ## Model definition parser.add_argument('--n_mels', type=int, default=80, help='Number of mel filterbanks'); parser.add_argument('--log_input', type=bool, default=False, help='Log input features') parser.add_argument('--model', type=str, default="", help='Name of model definition'); parser.add_argument('--encoder_type', type=str, default="SAP", help='Type of encoder'); parser.add_argument('--nOut', type=int, default=192, help='Embedding size in the last FC layer'); ## For test only parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only') ## Distributed and mixed precision training parser.add_argument('--port', type=str, default="7888", help='Port for distributed training, input as text'); parser.add_argument('--distributed', dest='distributed', action='store_true', help='Enable distributed training') parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training') args = parser.parse_args(); ## Parse YAML def find_option_type(key, parser): for opt in parser._get_optional_actions(): if ('--' + key) in opt.option_strings: return opt.type raise ValueError if args.config is not None: with open(args.config, "r") as f: yml_config = yaml.load(f, Loader=yaml.FullLoader) for k, v in yml_config.items(): if k in args.__dict__: typ = find_option_type(k, parser) args.__dict__[k] = typ(v) else: sys.stderr.write("Ignored unknown parameter {} in yaml.\n".format(k)) ## Try to import NSML try: import nsml from nsml import HAS_DATASET, DATASET_PATH, PARALLEL_WORLD, PARALLEL_PORTS, MY_RANK from nsml import NSML_NFS_OUTPUT, SESSION_NAME except: pass; warnings.simplefilter("ignore") ## ===== ===== ===== ===== ===== ===== ===== ===== ## Trainer script ## ===== ===== ===== ===== ===== ===== ===== ===== def LGL_threshold_update_gmm(loss_vals_path): with open(loss_vals_path, 'r') as f: lines = [line.strip().split() for line in f.readlines()] # losses = [float(line[0]) for line in lines] losses = [] errs = 0 for line in lines: try: losses.append(float(line[0])) except ValueError: errs += 1 pass if errs > 0: print('Could not read %d lines' % errs) log_losses = np.log(losses) gmm = GaussianMixture(n_components=2, random_state=0, covariance_type='full', tol=0.00001, max_iter=1000) gmm.fit(log_losses.reshape(-1, 1)) mean1 = gmm.means_[0, 0] covar1 = gmm.covariances_[0, 0] weight1 = gmm.weights_[0] x = np.linspace(min(log_losses), max(log_losses), 1000) g1 = weight1 * norm.pdf(x, mean1, np.sqrt(covar1)) mean2 = gmm.means_[1, 0] covar2 = gmm.covariances_[1, 0] weight2 = gmm.weights_[1] g2 = weight2 * norm.pdf(x, mean2, np.sqrt(covar2)) intersection = np.argwhere(np.diff(np.sign(g1 - g2))).flatten() max1 = x[np.argmax(g1)] max2 = x[np.argmax(g2)] good_intersection = x[intersection][(x[intersection] > min(max1, max2)) & (x[intersection] < max(max1, max2))] assert len(good_intersection) == 1, 'Wrong number of intersections' good_intersection = good_intersection[0] return good_intersection import idr_torch def main_worker(gpu, ngpus_per_node, args): args.gpu = gpu args.gpu = idr_torch.rank ngpus_per_node = idr_torch.size ## Load models s = SpeakerNet(**vars(args)); if args.distributed: # os.environ['MASTER_ADDR']='localhost' # os.environ['MASTER_PORT']=args.port # dist.init_process_group(backend='nccl', world_size=ngpus_per_node, rank=args.gpu, init_method='tcp://localhost:12345') dist.init_process_group(backend='nccl', world_size=ngpus_per_node, rank=args.gpu) torch.cuda.set_device(args.gpu) s.cuda(args.gpu) s = torch.nn.parallel.DistributedDataParallel(s, device_ids=[args.gpu])#, find_unused_parameters=True) print('Loaded the model on GPU {:d}'.format(args.gpu)) else: s = WrappedModel(s).cuda(args.gpu) it = 1 eers = [100]; if args.gpu == 0: ## Write args to scorefile scorefile = open(args.result_save_path+"/scores.txt", "a+"); ## Initialise trainer and data loader train_dataset = train_dataset_loader(**vars(args)) train_sampler = train_dataset_sampler(train_dataset, **vars(args)) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, num_workers=args.nDataLoaderThread, sampler=train_sampler, pin_memory=True, worker_init_fn=worker_init_fn, drop_last=True, ) # trainLoader = get_data_loader(args.train_list, **vars(args)); trainer = ModelTrainer(s, **vars(args)) ## Load model weights modelfiles = glob.glob('%s/model0*.model'%args.model_save_path) modelfiles.sort() if(args.initial_model != ""): trainer.loadParameters(args.initial_model); print("Model {} loaded!".format(args.initial_model)); elif len(modelfiles) >= 1: print("Model {} loaded from previous state!".format(modelfiles[-1])); trainer.loadParameters(modelfiles[-1]); it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1 for ii in range(1,it): trainer.__scheduler__.step() pytorch_total_params = sum(p.numel() for p in s.module.__S__.parameters()) print('Total parameters: ',pytorch_total_params) ## Evaluation code - must run on single GPU if args.eval == True: print('Test list',args.test_list) sc, lab, _, sc1,sc2 = trainer.evaluateFromList(**vars(args)) if args.gpu == 0: result = tuneThresholdfromScore(sc, lab, [1, 0.1]); result_s1 = tuneThresholdfromScore(sc1, lab, [1, 0.1]); result_s2 = tuneThresholdfromScore(sc2, lab, [1, 0.1]); fnrs, fprs, thresholds = ComputeErrorRates(sc, lab) mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, args.dcf_p_target, args.dcf_c_miss, args.dcf_c_fa) print('\n',time.strftime("%Y-%m-%d %H:%M:%S"), "VEER {:2.4f}".format(result[1]), "VEER_s1 {:2.4f}".format(result_s1[1]),"VEER_s2 {:2.4f}".format(result_s2[1]),"MinDCF {:2.5f}".format(mindcf)); if ("nsml" in sys.modules) and args.gpu == 0: training_report = {}; training_report["summary"] = True; training_report["epoch"] = it; training_report["step"] = it; training_report["val_eer"] = result[1]; training_report["val_dcf"] = mindcf; nsml.report(**training_report); return ## Save training code and params if args.gpu == 0: pyfiles = glob.glob('./*.py') strtime = datetime.datetime.now().strftime("%Y%m%d%H%M%S") zipf = zipfile.ZipFile(args.result_save_path+ '/run%s.zip'%strtime, 'w', zipfile.ZIP_DEFLATED) for file in pyfiles: zipf.write(file) zipf.close() with open(args.result_save_path + '/run%s.cmd'%strtime, 'w') as f: f.write('%s'%args) ## Core training script for it in range(it,args.max_epoch+1): train_sampler.set_epoch(it) clr = [x['lr'] for x in trainer.__optimizer__.param_groups] loss_vals_dir = 'exp/' + args.save_path.split('/')[-1] + '/loss_vals' os.makedirs(loss_vals_dir, exist_ok=True) loss_vals_path = os.path.join(loss_vals_dir, 'epoch_%d.txt' % it) if it >= 5: prev_loss_vals_path = os.path.join(loss_vals_dir, 'epoch_%d.txt' % (it - 1)) LGL_threshold = LGL_threshold_update_gmm(prev_loss_vals_path) # LGL_threshold = 1 if args.gpu == 0: if LGL_threshold is not None: print('Updated LGL threshold to %f' % LGL_threshold) else: print('Wrong number of intersections, keeping LGL threshold at %f' % LGL_threshold) trainer.update_lgl_threshold(LGL_threshold) loss, traineer = trainer.train_network(train_loader, loss_vals_path, it, verbose=(args.gpu == 0)) if args.distributed: dist.barrier() with open(loss_vals_path, 'w') as final_file: for r in range(dist.get_world_size()): part_file_path = f"{loss_vals_path.split('.')[0]}_rank{r}.txt" with open(part_file_path, 'r') as part_file: final_file.write(part_file.read()) if args.gpu == 0: print('\n',time.strftime("%Y-%m-%d %H:%M:%S"), "Epoch {:d}, TEER/TAcc {:2.2f}, TLOSS {:f}, LR {:f}".format(it, traineer.item(), loss.item(), max(clr))); scorefile.write("Epoch {:d}, TEER/TAcc {:2.2f}, TLOSS {:f}, LR {:f} \n".format(it, traineer.item(), loss.item(), max(clr))); if it % args.test_interval == 0: # sc, lab, _, as1, as2 = trainer.evaluateFromList(**vars(args)) if args.gpu == 0: trainer.saveParameters(args.model_save_path+"/model%09d.model"%it); scorefile.flush() if ("nsml" in sys.modules) and args.gpu == 0: training_report = {}; training_report["summary"] = True; training_report["epoch"] = it; training_report["step"] = it; training_report["train_loss"] = loss; training_report["min_eer"] = min(eers); nsml.report(**training_report); if args.gpu == 0: scorefile.close(); ## ===== ===== ===== ===== ===== ===== ===== ===== ## Main function ## ===== ===== ===== ===== ===== ===== ===== ===== def main(): # print(os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')) # os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' # print(os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')) if ("nsml" in sys.modules) and not args.eval: args.save_path = os.path.join(args.save_path,SESSION_NAME.replace('/','_')) args.model_save_path = args.save_path+"/model" args.result_save_path = args.save_path+"/result" args.feat_save_path = "" os.makedirs(args.model_save_path, exist_ok=True) os.makedirs(args.result_save_path, exist_ok=True) n_gpus = torch.cuda.device_count() print(n_gpus) print('Python Version:', sys.version) print('PyTorch Version:', torch.__version__) print('Number of GPUs:', torch.cuda.device_count()) print('Save path:',args.save_path) if args.distributed: # mp.spawn(main_worker, nprocs=n_gpus, args=(n_gpus, args)) main_worker(None, None, args) else: main_worker(0, None, args) if __name__ == '__main__': main()