trainEmbedNet.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. #!/usr/bin/python
  2. #-*- coding: utf-8 -*-
  3. import sys, time, os, argparse, socket
  4. import yaml
  5. import pdb
  6. import glob
  7. import datetime
  8. from utils import *
  9. from EmbedNet import *
  10. from DatasetLoader import get_data_loader
  11. import torchvision.transforms as transforms
  12. import mlflow
  13. # ## ===== ===== ===== ===== ===== ===== ===== =====
  14. # ## Parse arguments
  15. # ## ===== ===== ===== ===== ===== ===== ===== =====
  16. parser = argparse.ArgumentParser(description = "FaceNet");
  17. parser.add_argument('--config', type=str, default=None, help='Config YAML file');
  18. ## Data loader
  19. parser.add_argument('--batch_size', type=int, default=200, help='Batch size, number of classes per batch');
  20. parser.add_argument('--max_img_per_cls', type=int, default=500, help='Maximum number of images per class per epoch');
  21. parser.add_argument('--nDataLoaderThread', type=int, default=5, help='Number of loader threads');
  22. ## Training details
  23. parser.add_argument('--test_interval', type=int, default=5, help='Test and save every [test_interval] epochs');
  24. parser.add_argument('--max_epoch', type=int, default=100, help='Maximum number of epochs');
  25. parser.add_argument('--trainfunc', type=str, default="softmax", help='Loss function');
  26. ## Optimizer
  27. parser.add_argument('--optimizer', type=str, default="adam", help='sgd or adam');
  28. parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler');
  29. parser.add_argument('--lr', type=float, default=0.001, help='Learning rate');
  30. parser.add_argument("--lr_decay", type=float, default=0.90, help='Learning rate decay every [test_interval] epochs');
  31. parser.add_argument('--weight_decay', type=float, default=0, help='Weight decay in the optimizer');
  32. ## Loss functions
  33. parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions');
  34. parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions');
  35. parser.add_argument('--margin', type=float, default=0.1, help='Loss margin, only for some loss functions');
  36. parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions');
  37. parser.add_argument('--nPerClass', type=int, default=1, help='Number of images per class per batch, only for metric learning based losses');
  38. parser.add_argument('--nClasses', type=int, default=8700, help='Number of classes in the softmax layer, only for softmax-based losses');
  39. ## Load and save
  40. parser.add_argument('--initial_model', type=str, default="./models/amsoft_model.model", help='Initial model weights');
  41. parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs');
  42. ## Training and test data
  43. parser.add_argument('--train_path', type=str, default="data/vggface2", help='Absolute path to the train set');
  44. parser.add_argument('--train_ext', type=str, default="jpg", help='Training files extension');
  45. parser.add_argument('--test_path', type=str, default="data/test", help='Absolute path to the test set');
  46. parser.add_argument('--test_list', type=str, default="data/test_list.csv", help='Evaluation list');
  47. ## Model definition
  48. parser.add_argument('--model', type=str, default="ResNet18", help='Name of model definition');
  49. parser.add_argument('--nOut', type=int, default=512, help='Embedding size in the last FC layer');
  50. ## For test only
  51. parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only')
  52. ## Distributed and mixed precision training
  53. parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training')
  54. args = parser.parse_args();
  55. ## Parse YAML
  56. def find_option_type(key, parser):
  57. for opt in parser._get_optional_actions():
  58. if ('--' + key) in opt.option_strings:
  59. return opt.type
  60. raise ValueError
  61. if args.config is not None:
  62. with open(args.config, "r") as f:
  63. yml_config = yaml.load(f, Loader=yaml.FullLoader)
  64. for k, v in yml_config.items():
  65. if k in args.__dict__:
  66. typ = find_option_type(k, parser)
  67. args.__dict__[k] = typ(v)
  68. else:
  69. sys.stderr.write("Ignored unknown parameter {} in yaml.\n".format(k))
  70. # ## ===== ===== ===== ===== ===== ===== ===== =====
  71. # ## Trainer script
  72. # ## ===== ===== ===== ===== ===== ===== ===== =====
  73. def main_worker(args):
  74. ## mlflow
  75. with mlflow.start_run(run_name=args.lr):
  76. ## Load models
  77. s = EmbedNet(**vars(args)).cuda();
  78. it = 1
  79. ## Write args to scorefile
  80. scorefile = open(args.result_save_path+"/scores.txt", "a+");
  81. strtime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  82. scorefile.write('%s\n%s\n'%(strtime,args))
  83. scorefile.flush()
  84. ## Input transformations for training
  85. train_transform = transforms.Compose(
  86. [transforms.ToTensor(),
  87. transforms.Resize(256),
  88. transforms.RandomCrop([224,224]),
  89. transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
  90. ## Input transformations for evaluation
  91. test_transform = transforms.Compose(
  92. [transforms.ToTensor(),
  93. transforms.Resize(256),
  94. transforms.CenterCrop([224,224]),
  95. transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
  96. ## Initialise trainer and data loader
  97. trainLoader = get_data_loader(transform=train_transform, **vars(args));
  98. trainer = ModelTrainer(s, **vars(args))
  99. ## Load model weights
  100. modelfiles = glob.glob('%s/model0*.model'%args.model_save_path)
  101. modelfiles.sort()
  102. ## If the target directory already exists, start from the existing file
  103. if len(modelfiles) >= 1:
  104. trainer.loadParameters(modelfiles[-1]);
  105. print("Model %s loaded from previous state!"%modelfiles[-1]);
  106. it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1
  107. elif(args.initial_model != ""):
  108. trainer.loadParameters(args.initial_model);
  109. print("Model %s loaded!"%args.initial_model);
  110. ## If the current iteration is not 1, update the scheduler
  111. for ii in range(1,it):
  112. trainer.__scheduler__.step()
  113. ## Evaluation code
  114. if args.eval == True:
  115. sc, lab = trainer.evaluateFromList(transform=test_transform, **vars(args))
  116. result = tuneThresholdfromScore(sc, lab, [1, 0.1]);
  117. print('EER %2.4f'%(result[1]))
  118. quit();
  119. ## Core training script
  120. for it in range(it,args.max_epoch+1):
  121. clr = [x['lr'] for x in trainer.__optimizer__.param_groups]
  122. print(time.strftime("%Y-%m-%d %H:%M:%S"), it, "Training epoch %d with LR %f "%(it,max(clr)));
  123. loss, traineer = trainer.train_network(trainLoader, verbose=True);
  124. if it % args.test_interval == 0:
  125. sc, lab = trainer.evaluateFromList(transform=test_transform, **vars(args))
  126. result = tuneThresholdfromScore(sc, lab, [1, 0.1]);
  127. print("IT %d, VEER %2.4f"%(it, result[1]));
  128. scorefile.write("IT %d, VEER %2.4f\n"%(it, result[1]));
  129. ## mlflow logging
  130. mlflow.log_metric("VEER", float(result[1]))
  131. trainer.saveParameters(args.model_save_path+"/model%09d.model"%it);
  132. print(time.strftime("%Y-%m-%d %H:%M:%S"), "TEER/TAcc %2.2f, TLOSS %f"%( traineer, loss));
  133. scorefile.write("IT %d, TEER/TAcc %2.2f, TLOSS %f\n"%(it, traineer, loss));
  134. scorefile.flush()
  135. ## mlflow logging
  136. mlflow.log_param("lr", clr)
  137. mlflow.log_metric("TEER_TAcc", float(traineer))
  138. mlflow.log_metric("TLOSS", float(loss))
  139. mlflow.log_artifacts(args.result_save_path)
  140. scorefile.close();
  141. # ## ===== ===== ===== ===== ===== ===== ===== =====
  142. # ## Main function
  143. # ## ===== ===== ===== ===== ===== ===== ===== =====
  144. def main():
  145. args.model_save_path = args.save_path+"/model"
  146. args.result_save_path = args.save_path+"/result"
  147. if not(os.path.exists(args.model_save_path)):
  148. os.makedirs(args.model_save_path)
  149. if not(os.path.exists(args.result_save_path)):
  150. os.makedirs(args.result_save_path)
  151. main_worker(args)
  152. if __name__ == '__main__':
  153. main()