#!/usr/bin/python #-*- coding: utf-8 -*- import sys, time, os, argparse, socket import yaml import pdb import glob import datetime from utils import * from EmbedNet import * from DatasetLoader import get_data_loader import torchvision.transforms as transforms import mlflow # ## ===== ===== ===== ===== ===== ===== ===== ===== # ## Parse arguments # ## ===== ===== ===== ===== ===== ===== ===== ===== parser = argparse.ArgumentParser(description = "FaceNet"); parser.add_argument('--config', type=str, default=None, help='Config YAML file'); ## Data loader parser.add_argument('--batch_size', type=int, default=200, help='Batch size, number of classes per batch'); parser.add_argument('--max_img_per_cls', type=int, default=500, help='Maximum number of images per class per epoch'); parser.add_argument('--nDataLoaderThread', type=int, default=5, help='Number of loader threads'); ## Training details parser.add_argument('--test_interval', type=int, default=5, help='Test and save every [test_interval] epochs'); parser.add_argument('--max_epoch', type=int, default=100, help='Maximum number of epochs'); parser.add_argument('--trainfunc', type=str, default="softmax", help='Loss function'); ## Optimizer parser.add_argument('--optimizer', type=str, default="adam", 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.90, help='Learning rate decay every [test_interval] epochs'); parser.add_argument('--weight_decay', type=float, default=0, help='Weight decay in the optimizer'); ## 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.1, 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('--nPerClass', type=int, default=1, help='Number of images per class per batch, only for metric learning based losses'); parser.add_argument('--nClasses', type=int, default=8700, help='Number of classes in the softmax layer, only for softmax-based losses'); ## Load and save parser.add_argument('--initial_model', type=str, default="./models/amsoft_model.model", 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_path', type=str, default="data/vggface2", help='Absolute path to the train set'); parser.add_argument('--train_ext', type=str, default="jpg", help='Training files extension'); parser.add_argument('--test_path', type=str, default="data/test", help='Absolute path to the test set'); parser.add_argument('--test_list', type=str, default="data/test_list.csv", help='Evaluation list'); ## Model definition parser.add_argument('--model', type=str, default="ResNet18", help='Name of model definition'); parser.add_argument('--nOut', type=int, default=512, 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('--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)) # ## ===== ===== ===== ===== ===== ===== ===== ===== # ## Trainer script # ## ===== ===== ===== ===== ===== ===== ===== ===== def main_worker(args): ## mlflow with mlflow.start_run(run_name=args.lr): ## Load models s = EmbedNet(**vars(args)).cuda(); it = 1 ## Write args to scorefile scorefile = open(args.result_save_path+"/scores.txt", "a+"); strtime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") scorefile.write('%s\n%s\n'%(strtime,args)) scorefile.flush() ## Input transformations for training train_transform = transforms.Compose( [transforms.ToTensor(), transforms.Resize(256), transforms.RandomCrop([224,224]), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) ## Input transformations for evaluation test_transform = transforms.Compose( [transforms.ToTensor(), transforms.Resize(256), transforms.CenterCrop([224,224]), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) ## Initialise trainer and data loader trainLoader = get_data_loader(transform=train_transform, **vars(args)); trainer = ModelTrainer(s, **vars(args)) ## Load model weights modelfiles = glob.glob('%s/model0*.model'%args.model_save_path) modelfiles.sort() ## If the target directory already exists, start from the existing file if len(modelfiles) >= 1: trainer.loadParameters(modelfiles[-1]); print("Model %s loaded from previous state!"%modelfiles[-1]); it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1 elif(args.initial_model != ""): trainer.loadParameters(args.initial_model); print("Model %s loaded!"%args.initial_model); ## If the current iteration is not 1, update the scheduler for ii in range(1,it): trainer.__scheduler__.step() ## Evaluation code if args.eval == True: sc, lab = trainer.evaluateFromList(transform=test_transform, **vars(args)) result = tuneThresholdfromScore(sc, lab, [1, 0.1]); print('EER %2.4f'%(result[1])) quit(); ## Core training script for it in range(it,args.max_epoch+1): clr = [x['lr'] for x in trainer.__optimizer__.param_groups] print(time.strftime("%Y-%m-%d %H:%M:%S"), it, "Training epoch %d with LR %f "%(it,max(clr))); loss, traineer = trainer.train_network(trainLoader, verbose=True); if it % args.test_interval == 0: sc, lab = trainer.evaluateFromList(transform=test_transform, **vars(args)) result = tuneThresholdfromScore(sc, lab, [1, 0.1]); print("IT %d, VEER %2.4f"%(it, result[1])); scorefile.write("IT %d, VEER %2.4f\n"%(it, result[1])); ## mlflow logging mlflow.log_metric("VEER", float(result[1])) trainer.saveParameters(args.model_save_path+"/model%09d.model"%it); print(time.strftime("%Y-%m-%d %H:%M:%S"), "TEER/TAcc %2.2f, TLOSS %f"%( traineer, loss)); scorefile.write("IT %d, TEER/TAcc %2.2f, TLOSS %f\n"%(it, traineer, loss)); scorefile.flush() ## mlflow logging mlflow.log_param("lr", clr) mlflow.log_metric("TEER_TAcc", float(traineer)) mlflow.log_metric("TLOSS", float(loss)) mlflow.log_artifacts(args.result_save_path) scorefile.close(); # ## ===== ===== ===== ===== ===== ===== ===== ===== # ## Main function # ## ===== ===== ===== ===== ===== ===== ===== ===== def main(): args.model_save_path = args.save_path+"/model" args.result_save_path = args.save_path+"/result" if not(os.path.exists(args.model_save_path)): os.makedirs(args.model_save_path) if not(os.path.exists(args.result_save_path)): os.makedirs(args.result_save_path) main_worker(args) if __name__ == '__main__': main()