class ArcFace_Net(nn.Module): def __init__(self, backbone, test_model_path): super().__init__() # if backbone == 'resnet_face18': # self.model = resnet_face() # elif backbone == 'resnet18': # self.model = resnet18() # elif backbone == 'resnet34': # self.model = resnet34() # elif backbone == 'resnet50': # self.model = resnet50() # elif backbone == 'resnet101': # self.model = resnet101() # elif backbone == 'resnet152': # self.model = resnet152() if backbone == 'resnet50': self.model = Backbone(50, drop_ratio=0, mode='ir_se') elif backbone == 'resnet101': self.model = Backbone(100, drop_ratio=0, mode='ir_se') elif backbone == 'resnet152': self.model = Backbone(152, drop_ratio=0, mode='ir_se') if test_model_path: self.model.load_state_dict(torch.load(test_model_path)) print('arcface model loading...') def forward(self, x): self.model.eval() return self.model(x)
def __init__(self, use_mobilfacenet): self.embedding_size = 512 self.net_depth = 50 self.drop_ratio = 0.6 self.net_mode = 'ir_se' # or 'ir' self.threshold = 1.2 self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.face_detector = MTCNN() self.test_transform = trans.Compose([ trans.ToTensor(), trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ]) if use_mobilfacenet: self.model = MobileFaceNet(self.embedding_size).to(self.device) print('MobileFaceNet model generated') else: self.model = Backbone(self.net_depth, self.drop_ratio, self.net_mode).to(self.device) print('{}_{} model generated'.format(self.net_mode, self.net_depth)) self.threshold = self.threshold
def __init__(self, conf, inference=False): print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter( '/home/zzg/DeepLearning/InsightFace_Pytorch/work_space/log') self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 29000 #100 self.evaluate_every = len(self.loader) // 500 ##10 self.save_every = len(self.loader) // 290 #5 # self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) self.agedb_30, self.agedb_30_issame = get_val_data( '/home/zzg/DeepLearning/InsightFace_Pytorch/data/faces_emore') else: self.threshold = conf.threshold
def __init__(self, conf, inference=False): if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) # print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) # print('{}_{} model generated done !'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) # print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) rootdir = os.path.join(args.root_dir, args.rec_path) self.board_loss_every = len(self.loader) // len(self.loader) self.evaluate_every = len(self.loader) // 1 # self.save_every = len(self.loader)//len(self.loader) # 5 print('board loss every: {} -> evaluate_every: {} \n'.format( self.board_loss_every, self.evaluate_every)) print('loader paths of validation dataset {}'.format(rootdir)) self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( rootdir) else: self.threshold = conf.threshold
class face_learner(object): def __init__(self, use_mobilfacenet): self.embedding_size = 512 self.net_depth = 50 self.drop_ratio = 0.6 self.net_mode = 'ir_se' # or 'ir' self.threshold = 1.2 self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.face_detector = MTCNN() self.test_transform = trans.Compose([ trans.ToTensor(), trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ]) if use_mobilfacenet: self.model = MobileFaceNet(self.embedding_size).to(self.device) print('MobileFaceNet model generated') else: self.model = Backbone(self.net_depth, self.drop_ratio, self.net_mode).to(self.device) print('{}_{} model generated'.format(self.net_mode, self.net_depth)) self.threshold = self.threshold def load_state(self, modelFn): self.model.load_state_dict( torch.load(modelFn, map_location=torch.device(self.device))) def get_input_by_mtcnn(self, img): if (img.shape == (112, 112, 3)): pilImg = Image.fromarray(img) return pilImg else: bboxes, landmarks, faces = self.face_detector.align_multi( Image.fromarray(img), limit=None, min_face_size=60) if len(bboxes) < 1: return None idx = findBigFace(bboxes) return faces[idx] def get_feature(self, pilImg, tta=False): if tta: mirror = trans.functional.hflip(pilImg) emb = self.model( self.test_transform(pilImg).to(self.device).unsqueeze(0)) emb_mirror = self.model( self.test_transform(mirror).to(self.device).unsqueeze(0)) embeddings = l2_norm(emb + emb_mirror) else: embeddings = self.model( self.test_transform(pilImg).to(self.device).unsqueeze(0)) return embeddings
def __init__(self, conf, inference=False): print(conf) print(conf.use_mobilfacenet) input("CONF") if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones # Dataset Loader # ritorna un loader dataset ImageLoader self.loader, self.class_num = get_train_loader(conf) # Classe di tensorboardX per salvare i log # log_path indica il percorso dove sono salvate le statistiche self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') # paras_only_bn contiene i layer con i parametri della batchnorm paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') # Parametri che indicano ogni quanto salvare i modelli e le epoche self.board_loss_every = len(self.loader)//10 self.evaluate_every = len(self.loader)//10 self.save_every = len(self.loader)//5 print("DATASET") print(self.loader.dataset.root) # ritornano gli array e le labels delle diverse cartelle del dataset VALIDATION self.agedb_30 ,self.agedb_30_issame = get_val_data(self.loader.dataset.root.parent) else: self.threshold = conf.threshold
def __init__(self, conf, inference=False): print(conf) self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) print('loading... finish') if not inference: self.milestones = conf.milestones dataset = Dataset(root=conf.data_path, data_list_file=conf.datalist_path, phase='train', input_shape=(1, 112, 112)) self.loader = data.DataLoader(dataset, conf.batch_size) self.class_num = conf.class_num self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.load_state(conf, '2019-11-10-22-01.pth', True, False) # self.load_state(conf,'111.pth',False,True) print('load save') self.board_loss_every = len(self.loader) // 2 print(len(self.loader)) # self.board_loss_every = len(self.loader) // 100 self.evaluate_every = len(self.loader) // 2 # self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 2 # self.save_every = len(self.loader) // 5 # self.essex, self.essex_issame = get_val_data( # self.loader.dataset.root.parent) else: self.threshold = conf.threshold
def test(args): device = torch.device(('cuda:%d' % args.gpu) if torch.cuda.is_available() else 'cpu') BACKBONE = Backbone([args.input_size, args.input_size], args.num_layers, args.mode) BACKBONE.load_state_dict(torch.load(args.ckpt_path)) BACKBONE.to(device) BACKBONE.eval() transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) print('Start test at', datetime.now().strftime('%Y-%m-%d %H:%M:%S')) # accuracy with open(args.pair_file, 'r') as f: pairs = f.readlines() sims = [] labels = [] for pair_id, pair in tqdm.tqdm(enumerate(pairs)): # print('processing %d/%d...' % (pair_id, len(pairs)), end='\r') splits = pair.split() feat1 = get_feature(os.path.join(args.data_root, splits[0]), transform, BACKBONE, device) feat2 = get_feature(os.path.join(args.data_root, splits[1]), transform, BACKBONE, device) label = int(splits[2]) sim = np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2)) sims.append(sim) labels.append(label) acc, th = cal_accuracy(np.array(sims), np.array(labels)) print('acc=%f with threshold=%f' % (acc, th)) print('Finish test at', datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
def __init__(self, conf, inference=False, transfer=0, ext='final'): pprint.pprint(conf) self.conf = conf if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) tmp_idx = ext.rfind('_') # find the last '_' to replace it by '/' self.ext = '/' + ext[:tmp_idx] + '/' + ext[tmp_idx + 1:] self.writer = SummaryWriter(str(conf.log_path) + self.ext) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) self.optimizer = optim.Adam( list(self.model.parameters()) + list(self.head.parameters()), conf.lr) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.save_freq = len(self.loader) // 5 #//5 # originally, 100 self.evaluate_every = len(self.loader) #//5 # originally, 10 self.save_every = len(self.loader) #//2 # originally, 5 # self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) # self.val_112, self.val_112_issame = get_val_pair(self.loader.dataset.root.parent, 'val_112') else: self.threshold = conf.threshold self.train_losses = [] self.train_counter = [] self.test_losses = [] self.test_accuracy = [] self.test_counter = []
def __init__(self, conf, inference=False): print(conf) # self.loader, self.class_num = construct_msr_dataset(conf) self.loader, self.class_num = get_train_loader(conf) self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = QAMFace(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) self.focalLoss = FocalLoss() print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 1000 self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 2 else: self.threshold = conf.threshold # 多GPU训练 self.model = torch.nn.DataParallel(self.model) self.model.to(conf.device) self.head = torch.nn.DataParallel(self.head) self.head = self.head.to(conf.device)
def __init__(self, conf, inference=False): accuracy = 0.0 logger.debug(conf) if conf.use_mobilfacenet: # self.model = MobileFaceNet(conf.embedding_size).to(conf.device) self.model = MobileFaceNet(conf.embedding_size).cuda() logger.debug('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).cuda()#.to(conf.device) logger.debug('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones logger.info('loading data...') self.loader, self.class_num = get_train_loader(conf, 'emore', sample_identity=True) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = CircleLoss(m=0.25, gamma=256.0).cuda() paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn, 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) # self.optimizer = torch.nn.parallel.DistributedDataParallel(optimizer,device_ids=[conf.argsed]) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) if conf.fp16: self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level="O2") self.model = DistributedDataParallel(self.model).cuda() else: self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[conf.argsed]).cuda() #add line for distributed self.board_loss_every = len(self.loader)//100 self.evaluate_every = len(self.loader)//2 self.save_every = len(self.loader)//2 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(Path(self.loader.dataset.root).parent) else: self.threshold = conf.threshold self.loader, self.query_ds, self.gallery_ds = get_test_loader(conf)
def create_model(num_classes=21, device=torch.device('cpu')): backbone = Backbone(pretrain_path='./pretrain/resnet50.pth') model = SSD300(backbone=backbone, num_classes=num_classes) pre_ssd_path = './pretrain/nvidia_ssdpyt_fp32.pt' pre_model_dict = torch.load(pre_ssd_path, map_location=device) pre_weights_dict = pre_model_dict['model'] # only use the pre_trained bounding boxes regression weights del_conf_loc_dict = {} for k, v in pre_weights_dict.items(): split_key = k.split('.') if 'conf' in split_key: continue del_conf_loc_dict.update({k: v}) missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict, strict=False) # if len(missing_keys) != 0 or len(unexpected_keys) != 0: # print('missing_keys: ', missing_keys) # print('unexpected_keys: ', unexpected_keys) return model
def __init__(self, backbone, test_model_path): super().__init__() # if backbone == 'resnet_face18': # self.model = resnet_face() # elif backbone == 'resnet18': # self.model = resnet18() # elif backbone == 'resnet34': # self.model = resnet34() # elif backbone == 'resnet50': # self.model = resnet50() # elif backbone == 'resnet101': # self.model = resnet101() # elif backbone == 'resnet152': # self.model = resnet152() if backbone == 'resnet50': self.model = Backbone(50, drop_ratio=0, mode='ir_se') elif backbone == 'resnet101': self.model = Backbone(100, drop_ratio=0, mode='ir_se') elif backbone == 'resnet152': self.model = Backbone(152, drop_ratio=0, mode='ir_se') if test_model_path: self.model.load_state_dict(torch.load(test_model_path)) print('arcface model loading...')
def __init__(self, conf, inference=False): print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) print('class_num:', self.class_num) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') # if conf.data_mode == 'small_vgg': # self.board_loss_every = len(self.loader) # print('len(loader', len(self.loader)) # self.evaluate_every = len(self.loader) # self.save_every = len(self.loader) # # self.lfw, self.lfw_issame = get_val_data(conf, conf.smallvgg_folder) # else: # self.board_loss_every = len(self.loader) # self.evaluate_every = len(self.loader)//10 # self.save_every = len(self.loader)//5 self.agedb_30, self.cfp_fp, self.lfw, self.kface, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.kface_issame = get_val_data( conf, self.loader.dataset.root.parent) else: self.threshold = conf.threshold
class face_learner(object): def __init__(self, conf, inference=False): print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) print('class_num:', self.class_num) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') # if conf.data_mode == 'small_vgg': # self.board_loss_every = len(self.loader) # print('len(loader', len(self.loader)) # self.evaluate_every = len(self.loader) # self.save_every = len(self.loader) # # self.lfw, self.lfw_issame = get_val_data(conf, conf.smallvgg_folder) # else: # self.board_loss_every = len(self.loader) # self.evaluate_every = len(self.loader)//10 # self.save_every = len(self.loader)//5 self.agedb_30, self.cfp_fp, self.lfw, self.kface, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.kface_issame = get_val_data( conf, self.loader.dataset.root.parent) else: self.threshold = conf.threshold def save_state(self, conf, accuracy, e, loss, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path model_name = f'model_e:{e+1}_acc:{accuracy}_loss:{loss}_{extra}.pth' torch.save(self.model.state_dict(), save_path / model_name) if not model_only: # 똥째로 저장 state = { 'model': self.model.state_dict(), 'head': self.head.state_dict(), 'optimizer': self.optimizer.state_dict() } torch.save(state, save_path / model_name) # print('model saved: ', model_name) ## 따로따로 저장 # torch.save( # self.head.state_dict(), save_path / # ('head_{}_accuracy:{}_epoch:{}_step:{}_{}.pth'.format(get_time(), accuracy, e, self.step, extra))) # torch.save( # self.optimizer.state_dict(), save_path / # ('optimizer_{}_accuracy:{}_epoch:{}_step:{}_{}.pth'.format(get_time(), accuracy, e, self.step, extra))) def load_state(self, conf, fixed_str, from_save_folder=False, model_only=False): if from_save_folder: save_path = conf.save_path else: save_path = conf.model_path if model_only: self.model.load_state_dict( torch.load(save_path / 'model_{}'.format(fixed_str))) if not model_only: model_name = 'model_{}'.format(fixed_str) state = torch.load(save_path / model_name) self.model.load_state_dict(state['model']) self.head.load_state_dict(state['head']) self.optimizer.load_state_dict(state['optimizer']) # self.head.load_state_dict(torch.load(save_path/'head_{}'.format(fixed_str))) # self.optimizer.load_state_dict(torch.load(save_path/'optimizer_{}'.format(fixed_str))) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) # self.writer.add_scalar('{}_val:true accept ratio'.format(db_name), val, self.step) # self.writer.add_scalar('{}_val_std'.format(db_name), val_std, self.step) # self.writer.add_scalar('{}_far:False Acceptance Ratio'.format(db_name), far, self.step) def evaluate(self, conf, carray, issame, nrof_folds=5, tta=False): self.model.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size]) with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch) else: embeddings[idx:idx + conf.batch_size] = self.model( batch.to(conf.device)).cpu() idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:] = l2_norm(emb_batch) else: embeddings[idx:] = self.model(batch.to(conf.device)).cpu() tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds) buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor def find_lr(self, conf, init_value=1e-8, final_value=10., beta=0.98, bloding_scale=3., num=None): if not num: num = len(self.loader) mult = (final_value / init_value)**(1 / num) lr = init_value for params in self.optimizer.param_groups: params['lr'] = lr self.model.train() avg_loss = 0. best_loss = 0. batch_num = 0 losses = [] log_lrs = [] for i, (imgs, labels) in tqdm(enumerate(self.loader), total=num): imgs = imgs.to(conf.device) labels = labels.to(conf.device) batch_num += 1 self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) #Compute the smoothed loss avg_loss = beta * avg_loss + (1 - beta) * loss.item() self.writer.add_scalar('avg_loss', avg_loss, batch_num) smoothed_loss = avg_loss / (1 - beta**batch_num) self.writer.add_scalar('smoothed_loss', smoothed_loss, batch_num) #Stop if the loss is exploding if batch_num > 1 and smoothed_loss > bloding_scale * best_loss: print('exited with best_loss at {}'.format(best_loss)) plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses #Record the best loss if smoothed_loss < best_loss or batch_num == 1: best_loss = smoothed_loss #Store the values losses.append(smoothed_loss) log_lrs.append(math.log10(lr)) self.writer.add_scalar('log_lr', math.log10(lr), batch_num) #Do the SGD step #Update the lr for the next step loss.backward() self.optimizer.step() lr *= mult for params in self.optimizer.param_groups: params['lr'] = lr if batch_num > num: plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses def train(self, conf, epochs): self.model.train() running_loss = 0. time_ = datetime.datetime.now() # check parameter of model print("------------------------------------------------------------") total_params = sum(p.numel() for p in self.model.parameters()) print("num of parameter : ", total_params) trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) print("num of trainable_ parameter :", trainable_params) print("------------------------------------------------------------") # if conf.data_mode == 'small_vgg': for e in range(epochs): print('epoch {} started'.format(e)) if e == self.milestones[0]: self.schedule_lr() if e == self.milestones[1]: self.schedule_lr() if e == self.milestones[2]: self.schedule_lr() accuracy_list = [] for iter_, (imgs, labels) in tqdm(enumerate(iter(self.loader))): # print('iter_', type(iter_)) # print('step', self.step) imgs = imgs.to(conf.device) labels = labels.to(conf.device) self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) loss.backward() running_loss += loss.item() self.optimizer.step() if iter_ % conf.print_iter == 0: elapsed = datetime.datetime.now() - time_ expected = elapsed * (conf.batch_size / conf.print_iter) _epoch = round(e + ((iter_ + 1) / conf.batch_size), 2) # print('_epoch', _epoch) _loss = round(loss.item(), 5) print( f'[{_epoch}/{conf.epochs}] loss:{_loss}, elapsed:{elapsed}, expected per epoch: {expected}' ) time_ = datetime.datetime.now() self.step += 1 # log 남기기 board_loss_every = len(self.loader) / conf.print_iter print('board_loss_ebery', board_loss_every) loss_board = running_loss / board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. # validate accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.agedb_30, self.agedb_30_issame) self.board_val('agedb_30', accuracy, best_threshold, roc_curve_tensor) accuracy_list.append(accuracy) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.lfw, self.lfw_issame) self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor) accuracy_list.append(accuracy) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.cfp_fp, self.cfp_fp_issame) self.board_val('cfp_fp', accuracy, best_threshold, roc_curve_tensor) accuracy_list.append(accuracy) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.cfp_fp, self.cfp_fp_issame) self.board_val('kface', accuracy, best_threshold, roc_curve_tensor) accuracy_list.append(accuracy) # save model, info # print(accuracy_list) accuracy_mean = round(sum(accuracy_list) / conf.testset_num, 5) loss_mean = round(loss_board, 5) self.save_state(conf, accuracy_mean, e, loss_mean, to_save_folder=False, extra='training') time_ = datetime.datetime.now() elapsed = datetime.datetime.now() - time_ print( f'[epoch {e + 1}] acc: {accuracy_mean}, loss: {loss_mean}, elapsed: {elapsed}' ) # print('train_loss:', loss_board) self.save_state(conf, accuracy_mean, e, loss_mean, to_save_folder=False, extra='final') def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print('***self.optimizer:', self.optimizer) def infer(self, conf, faces, target_embs, tta=False): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model( conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model( conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append( self.model( conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose( 1, 0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum
def __init__(self, conf, inference=False, transfer=0, ext='final'): pprint.pprint(conf) self.conf = conf if conf.arch == "mobile": self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') elif conf.arch == "ir_se": self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.arch).to(conf.device) print('{}_{} model generated'.format(conf.arch, conf.net_depth)) elif conf.arch == "resnet50": self.model = ResNet(embedding_size=512, arch=conf.arch).to(conf.device) print("resnet model {} generated".format(conf.arch)) else: exit("model not supported yet!") if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) tmp_idx = ext.rfind('_') # find the last '_' to replace it by '/' self.ext = '/' + ext[:tmp_idx] + '/' + ext[tmp_idx + 1:] self.writer = SummaryWriter(str(conf.log_path) + self.ext) self.step = 0 print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if transfer == 3: self.optimizer = optim.Adam( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr) # , momentum = conf.momentum) elif transfer == 2: self.optimizer = optim.Adam( [ { 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr) # , momentum = conf.momentum) elif transfer == 1: self.optimizer = optim.Adam( [ { 'params': [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr) # , momentum = conf.momentum) else: """ self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) """ self.optimizer = optim.Adam(list(self.model.parameters()) + list(self.head.parameters()), lr=conf.lr) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.save_freq = len(self.loader) #//5 # originally, 100 self.evaluate_every = len(self.loader) #//5 # originally, 10 self.save_every = len(self.loader) #//2 # originally, 5 # self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) # self.val_112, self.val_112_issame = get_val_pair(self.loader.dataset.root.parent, 'val_112') else: self.threshold = conf.threshold self.train_losses = [] self.train_counter = [] self.test_losses = [] self.test_accuracy = [] self.test_counter = []
class face_learner(object): def __init__(self, conf, inference=False): print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader)//100 self.evaluate_every = len(self.loader)//10 self.save_every = len(self.loader)//5 self.lfw, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) else: self.threshold = conf.threshold def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.model.state_dict(), save_path / ('model_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.head.state_dict(), save_path / ('head_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) def load_state(self, conf, fixed_str, from_save_folder=False, model_only=False): if from_save_folder: save_path = conf.save_path else: save_path = conf.model_path self.model.load_state_dict(torch.load(save_path/'model_{}'.format(fixed_str))) if not model_only: self.head.load_state_dict(torch.load(save_path/'head_{}'.format(fixed_str))) self.optimizer.load_state_dict(torch.load(save_path/'optimizer_{}'.format(fixed_str))) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) # self.writer.add_scalar('{}_val:true accept ratio'.format(db_name), val, self.step) # self.writer.add_scalar('{}_val_std'.format(db_name), val_std, self.step) # self.writer.add_scalar('{}_far:False Acceptance Ratio'.format(db_name), far, self.step) def evaluate(self, conf, carray, issame, nrof_folds = 5, tta = False): self.model.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size]) with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device)) embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch) else: embeddings[idx:idx + conf.batch_size] = self.model(batch.to(conf.device)).cpu() idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device)) embeddings[idx:] = l2_norm(emb_batch) else: embeddings[idx:] = self.model(batch.to(conf.device)).cpu() tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds) buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor def train(self, conf, epochs): self.model.train() running_loss = 0. for e in range(epochs): print('epoch {} started'.format(e)) if e == self.milestones[0]: self.schedule_lr() if e == self.milestones[1]: self.schedule_lr() if e == self.milestones[2]: self.schedule_lr() for imgs, labels in tqdm(iter(self.loader)): imgs = imgs.to(conf.device) labels = labels.to(conf.device) self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) loss.backward() running_loss += loss.item() self.optimizer.step() if self.step % (self.board_loss_every+1) == 0 and self.step != 0: loss_board = running_loss / (self.board_loss_every+1) self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % (self.evaluate_every+1) == 0 and self.step != 0: accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.lfw, self.lfw_issame) self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor) self.model.train() if self.step % (self.save_every+1) == 0 and self.step != 0: self.save_state(conf, accuracy) self.step += 1 self.save_state(conf, accuracy, to_save_folder=True, extra='final') def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer) def infer(self, conf, faces, target_embs, tta=False): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model(conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model(conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append(self.model(conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose(1,0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum
class face_learner(object): def __init__(self, conf, inference=False): print(conf) self.lr=conf.lr if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: ############################### ir_se50 ######################################## if conf.struct =='ir_se_50': self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) ############################### resnet101 ###################################### if conf.struct =='ir_se_101': self.model = resnet101().to(conf.device) print('resnet101 model generated') if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 ############################### ir_se50 ######################################## if conf.struct =='ir_se_50': self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) self.head_race = Arcface(embedding_size=conf.embedding_size, classnum=4).to(conf.device) ############################### resnet101 ###################################### if conf.struct =='ir_se_101': self.head = ArcMarginModel(embedding_size=conf.embedding_size,classnum=self.class_num).to(conf.device) self.head_race = ArcMarginModel(embedding_size=conf.embedding_size,classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel] + [self.head_race.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn + [self.head.kernel] + [self.head_race.kernel], 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') print('len of loader:',len(self.loader)) self.board_loss_every = len(self.loader)//min(len(self.loader),100) self.evaluate_every = len(self.loader)//1 self.save_every = len(self.loader)//1 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(conf.val_folder) else: #self.threshold = conf.threshold pass def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.model.state_dict(), save_path / ('model_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.head.state_dict(), save_path / ('head_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) torch.save( self.head_race.state_dict(), save_path / ('head__race{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{}_step:{}_{}.pth'.format(get_time(), accuracy, self.step, extra))) def load_state(self, model, head=None,head_race=None,optimizer=None): self.model.load_state_dict(torch.load(model),strict=False) if head is not None: self.head.load_state_dict(torch.load(head)) if head_race is not None: self.head_race.load_state_dict(torch.load(head_race)) if optimizer is not None: self.optimizer.load_state_dict(torch.load(optimizer)) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor,tpr_val): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) self.writer.add_scalar('{}[email protected]'.format(db_name), tpr_val, self.step) # self.writer.add_scalar('{}_val:true accept ratio'.format(db_name), val, self.step) # self.writer.add_scalar('{}_val_std'.format(db_name), val_std, self.step) # self.writer.add_scalar('{}_far:False Acceptance Ratio'.format(db_name), far, self.step) def evaluate(self, conf, carray, issame, nrof_folds = 5, tta = False): self.model.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size]) with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device)) embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch) else: embeddings[idx:idx + conf.batch_size] = self.model(batch.to(conf.device)).cpu() idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model(fliped.to(conf.device)) embeddings[idx:] = l2_norm(emb_batch) else: embeddings[idx:] = self.model(batch.to(conf.device)).cpu() tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds) try: tpr_val = tpr[np.less(fpr,0.0012)&np.greater(fpr,0.0008)][0] except: tpr_val = 0 buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor,tpr_val def find_lr(self, conf, init_value=1e-8, final_value=10., beta=0.98, bloding_scale=3., num=None): if not num: num = len(self.loader) mult = (final_value / init_value)**(1 / num) lr = init_value for params in self.optimizer.param_groups: params['lr'] = lr self.model.train() avg_loss = 0. best_loss = 0. batch_num = 0 losses = [] log_lrs = [] for i, (imgs, labels) in tqdm(enumerate(self.loader), total=num): imgs = imgs.to(conf.device) labels = labels.to(conf.device) batch_num += 1 self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) #Compute the smoothed loss avg_loss = beta * avg_loss + (1 - beta) * loss.item() self.writer.add_scalar('avg_loss', avg_loss, batch_num) smoothed_loss = avg_loss / (1 - beta**batch_num) self.writer.add_scalar('smoothed_loss', smoothed_loss,batch_num) #Stop if the loss is exploding if batch_num > 1 and smoothed_loss > bloding_scale * best_loss: print('exited with best_loss at {}'.format(best_loss)) plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses #Record the best loss if smoothed_loss < best_loss or batch_num == 1: best_loss = smoothed_loss #Store the values losses.append(smoothed_loss) log_lrs.append(math.log10(lr)) self.writer.add_scalar('log_lr', math.log10(lr), batch_num) #Do the SGD step #Update the lr for the next step loss.backward() self.optimizer.step() lr *= mult for params in self.optimizer.param_groups: params['lr'] = lr if batch_num > num: plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses def train(self, conf, epochs): self.model = self.model.to(conf.device) self.head = self.head.to(conf.device) self.head_race = self.head_race.to(conf.device) self.model.train() self.head.train() self.head_race.train() running_loss = 0. for e in range(epochs): print('epoch {} started'.format(e)) if e == 8:#5 #train hear_race #self.init_lr() conf.loss0 = False conf.loss1 = True conf.loss2 = True conf.model = False conf.head = False conf.head_race = True print(conf) if e == 16:#10: #self.init_lr() self.schedule_lr() conf.loss0 = True conf.loss1 = True conf.loss2 = True conf.model = True conf.head = True conf.head_race = True print(conf) if e == 28:#22 self.schedule_lr() if e == 32: self.schedule_lr() if e == 35: self.schedule_lr() requires_grad(self.head,conf.head) requires_grad(self.head_race,conf.head_race) requires_grad(self.model,conf.model) for imgs, labels in tqdm(iter(self.loader)): imgs = imgs.to(conf.device) labels = labels.to(conf.device) labels_race = torch.zeros_like(labels) race0_index = labels.lt(sum(conf.race_num[:1])) race1_index = labels.lt(sum(conf.race_num[:2])) & labels.ge(sum(conf.race_num[:1])) race2_index = labels.lt(sum(conf.race_num[:3])) & labels.ge(sum(conf.race_num[:2])) race3_index = labels.ge(sum(conf.race_num[:3])) labels_race[race0_index]=0 labels_race[race1_index] = 1 labels_race[race2_index] = 2 labels_race[race3_index] = 3 self.optimizer.zero_grad() embeddings = self.model(imgs) thetas ,w = self.head(embeddings, labels) thetas_race ,w_race = self.head_race(embeddings, labels_race) loss = 0 loss0 = conf.ce_loss(thetas, labels) loss1 = conf.ce_loss(thetas_race, labels_race) loss2 = torch.mm(w_race.t(),w).to(conf.device) target = torch.zeros_like(loss2).to(conf.device) target[0][:sum(conf.race_num[:1])] = 1 target[1][sum(conf.race_num[:1]):sum(conf.race_num[:2])] = 1 target[2][sum(conf.race_num[:2]):sum(conf.race_num[:3])] = 1 target[3][sum(conf.race_num[:3]):] = 1 weight = torch.zeros_like(loss2).to(conf.device) for i in range(4): weight[i,:] = sum(conf.race_num)/conf.race_num[i] #loss2 = torch.nn.functional.mse_loss(loss2 , target) loss2 = F.binary_cross_entropy(torch.sigmoid(loss2),target,weight) if conf.loss0 ==True: loss += 2*loss0 if conf.loss1 ==True: loss += loss1 if conf.loss2 ==True: loss += loss2 #loss = loss0 + loss1 + loss2 loss.backward() running_loss += loss.item() self.optimizer.step() if self.step % self.board_loss_every == 0 and self.step != 0: loss_board = running_loss / self.board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % self.evaluate_every == 0 and self.step != 0: accuracy=None accuracy, best_threshold, roc_curve_tensor ,tpr_val= self.evaluate(conf, self.agedb_30, self.agedb_30_issame) self.board_val('agedb_30', accuracy, best_threshold, roc_curve_tensor,tpr_val) accuracy, best_threshold, roc_curve_tensor,tpr_val = self.evaluate(conf, self.lfw, self.lfw_issame) self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor,tpr_val) accuracy, best_threshold, roc_curve_tensor,tpr_val = self.evaluate(conf, self.cfp_fp, self.cfp_fp_issame) self.board_val('cfp_fp', accuracy, best_threshold, roc_curve_tensor,tpr_val) self.model.train() if self.step % self.save_every == 0 and self.step != 0: self.save_state(conf, accuracy) self.step += 1 self.save_state(conf, accuracy, to_save_folder=True, extra='final') def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer) def init_lr(self): for params in self.optimizer.param_groups: params['lr'] = self.lr print(self.optimizer) def schedule_lr_add(self): for params in self.optimizer.param_groups: params['lr'] *= 10 print(self.optimizer) def infer(self, conf, faces, target_embs, tta=False): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model(conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model(conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append(self.model(conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose(1,0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum
def test_Backbone_output_shape(): bb = Backbone(2, 3, 100) input = torch.rand(1, 3, 16, 16) assert bb(input).shape == (1, 2, 100, 2, 2)
def __init__(self, conf, inference=False): print(conf) self.lr=conf.lr if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: ############################### ir_se50 ######################################## if conf.struct =='ir_se_50': self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) ############################### resnet101 ###################################### if conf.struct =='ir_se_101': self.model = resnet101().to(conf.device) print('resnet101 model generated') if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 ############################### ir_se50 ######################################## if conf.struct =='ir_se_50': self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) self.head_race = Arcface(embedding_size=conf.embedding_size, classnum=4).to(conf.device) ############################### resnet101 ###################################### if conf.struct =='ir_se_101': self.head = ArcMarginModel(embedding_size=conf.embedding_size,classnum=self.class_num).to(conf.device) self.head_race = ArcMarginModel(embedding_size=conf.embedding_size,classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.kernel] + [self.head_race.kernel], 'weight_decay': 4e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) else: self.optimizer = optim.SGD([ {'params': paras_wo_bn + [self.head.kernel] + [self.head_race.kernel], 'weight_decay': 5e-4}, {'params': paras_only_bn} ], lr = conf.lr, momentum = conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') print('len of loader:',len(self.loader)) self.board_loss_every = len(self.loader)//min(len(self.loader),100) self.evaluate_every = len(self.loader)//1 self.save_every = len(self.loader)//1 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(conf.val_folder) else: #self.threshold = conf.threshold pass
def __init__(self, conf, inference=False): self.backbone = Backbone().to(conf.device) self.idprehead = PreheadID().to(conf.device) self.idhead = Arcface().to(conf.device) self.attrhead = Attrhead().to(conf.device) print('model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones train_dataset = CelebA( 'dataset', 'celebA_train.txt', trans.Compose([ trans.RandomHorizontalFlip(), trans.ToTensor(), trans.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])) valid_dataset = CelebA( 'dataset', 'celebA_validation.txt', trans.Compose([ trans.RandomHorizontalFlip(), trans.ToTensor(), trans.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])) self.loader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, pin_memory=conf.pin_memory, num_workers=conf.num_workers) self.valid_loader = DataLoader(valid_dataset, batch_size=conf.batch_size, shuffle=True, pin_memory=conf.pin_memory, num_workers=conf.num_workers) self.writer = SummaryWriter(conf.log_path) self.step = 0 paras_only_bn_1, paras_wo_bn_1 = separate_bn_paras(self.backbone) paras_only_bn_2, paras_wo_bn_2 = separate_bn_paras(self.idprehead) paras_only_bn_3, paras_wo_bn_3 = separate_bn_paras(self.attrhead) paras_only_bn = paras_only_bn_1 + paras_only_bn_2 + paras_only_bn_3 paras_wo_bn = paras_wo_bn_1 + paras_wo_bn_2 + paras_wo_bn_3 self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.idhead.kernel], 'weight_decay': 1e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 8 self.evaluate_every = len(self.loader) // 4 self.save_every = len(self.loader) // 2 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( Path("data/faces_emore/")) else: self.threshold = conf.threshold
def __init__(self, conf, inference=False, transfer=0): pprint.pprint(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: if transfer == 3: self.optimizer = optim.SGD( [{ 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) elif transfer == 2: self.optimizer = optim.SGD([ { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr, momentum=conf.momentum) elif transfer == 1: self.optimizer = optim.SGD([ { 'params': [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 5 # originally, 100 self.evaluate_every = len(self.loader) // 5 # originally, 10 self.save_every = len(self.loader) // 2 # originally, 5 # self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) self.val_112, self.val_112_issame = get_val_pair( self.loader.dataset.root.parent, 'val_112') else: self.threshold = conf.threshold
def train(args): DEVICE = torch.device(("cuda:%d"%args.gpu[0]) if torch.cuda.is_available() else "cpu") writer = SummaryWriter(args.log_root) train_transform = transforms.Compose([transforms.Resize([int(128*args.input_size/112), int(128*args.input_size/112)]), transforms.RandomCrop([args.input_size, args.input_size]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[args.rgb_mean,args.rgb_mean,args.rgb_mean], std=[args.rgb_std,args.rgb_std,args.rgb_std]) ]) train_dataset = datasets.ImageFolder(args.data_root, train_transform) weights = make_weights_for_balanced_classes(train_dataset.imgs, len(train_dataset.classes)) weights = torch.DoubleTensor(weights) sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights)) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=8, shuffle=True, drop_last=True) NUM_CLASS = len(train_loader.dataset.classes) BACKBONE = Backbone([args.input_size, args.input_size], args.num_layers, args.mode) HEAD = ArcFace(args.emb_dims, NUM_CLASS, device_id=args.gpu) LOSS = FocalLoss() backbone_paras_only_bn, backbone_paras_wo_bn = separate_irse_bn_paras(BACKBONE) _, head_paras_wo_bn = separate_irse_bn_paras(HEAD) optimizer = optim.SGD([{'params': backbone_paras_wo_bn+head_paras_wo_bn, 'weight_decay': args.weight_decay}, {'params': backbone_paras_only_bn}], lr=args.lr, momentum=args.momentum) # optimizer = optim.AdamW([{'params': backbone_paras_wo_bn+head_paras_wo_bn, 'weight_decay': args.weight_decay}, # {'params': backbone_paras_only_bn}], lr=args.lr, momentum=args.momentum) if args.load_ckpt: BACKBONE.load_state_dict(torch.load(os.path.join(args.load_ckpt, 'backbone_epoch{}.pth'.format(args.load_epoch)))) HEAD.load_state_dict(torch.load(os.path.join(args.load_ckpt, 'head_epoch{}.pth'.format(args.load_epoch)))) print('Checkpoint loaded') start_epoch = args.load_epoch if args.load_ckpt else 0 BACKBONE = nn.DataParallel(BACKBONE, device_ids=args.gpu) BACKBONE = BACKBONE.to(DEVICE) dispaly_frequency = len(train_loader) // 100 NUM_EPOCH_WARM_UP = args.num_epoch // 25 NUM_BATCH_WARM_UP = len(train_loader) * NUM_EPOCH_WARM_UP batch = 0 print('Start training at %s!' % datetime.now().strftime('%Y-%m-%d %H:%M:%S')) for epoch in range(start_epoch, args.num_epoch): if epoch==args.stages[0] or epoch==args.stages[1] or epoch==args.stages[2]: for params in optimizer.param_groups: params['lr'] /= 10. BACKBONE.train() HEAD.train() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() for inputs, labels in train_loader: if (epoch+1 <= NUM_EPOCH_WARM_UP) and (batch+1 <= NUM_BATCH_WARM_UP): for params in optimizer.param_groups: params['lr'] = (batch+1) * args.lr / NUM_BATCH_WARM_UP inputs = inputs.to(DEVICE) labels = labels.to(DEVICE).long() features = BACKBONE(inputs) outputs = HEAD(features, labels) loss = LOSS(outputs, labels) prec1, prec5 = accuracy(outputs.data, labels, topk=(1,5)) losses.update(loss.data.item(), inputs.size(0)) top1.update(prec1.data.item(), inputs.size(0)) top5.update(prec5.data.item(), inputs.size(0)) optimizer.zero_grad() loss.backward() optimizer.step() batch += 1 if batch % dispaly_frequency == 0: print('%s Epoch %d/%d Batch %d/%d: train loss %f, train prec@1 %f, train prec@5 %f' % (datetime.now().strftime('%Y-%m-%d %H:%M:%S'), epoch, args.num_epoch, batch, len(train_loader)*args.num_epoch, losses.avg, top1.avg, top5.avg)) writer.add_scalar('Train_Loss', losses.avg, epoch+1) writer.add_scalar('Train_Top1_Accuracy', top1.avg, epoch+1) writer.add_scalar('Train_Top5_Accuracy', top5.avg, epoch+1) torch.save(BACKBONE.module.state_dict(), os.path.join(args.ckpt_root, 'backbone_epoch%d.pth'%(epoch+1))) torch.save(HEAD.state_dict(), os.path.join(args.ckpt_root, 'head_epoch%d.pth'%(epoch+1)))
class face_learner(object): def __init__(self, conf, inference=False): print(conf) # self.loader, self.class_num = construct_msr_dataset(conf) self.loader, self.class_num = get_train_loader(conf) self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = QAMFace(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) self.focalLoss = FocalLoss() print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 1000 self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 2 else: self.threshold = conf.threshold # 多GPU训练 self.model = torch.nn.DataParallel(self.model) self.model.to(conf.device) self.head = torch.nn.DataParallel(self.head) self.head = self.head.to(conf.device) def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.model.state_dict(), save_path / ('model_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.head.state_dict(), save_path / ('head_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) def load_state(self, conf, fixed_str, from_save_folder=False, model_only=False): print('resume model from ' + fixed_str) if from_save_folder: save_path = conf.save_path else: save_path = conf.model_path self.model.load_state_dict( torch.load(save_path / 'model_{}'.format(fixed_str))) if not model_only: self.head.load_state_dict( torch.load(save_path / 'head_{}'.format(fixed_str))) self.optimizer.load_state_dict( torch.load(save_path / 'optimizer_{}'.format(fixed_str))) def board_val(self, db_name, accuracy, best_threshold=0, roc_curve_tensor=0): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) def train(self, conf, epochs): self.model.train() running_loss = 0. for e in range(epochs): print('epoch {} started'.format(e)) # manually decay lr if e in self.milestones: self.schedule_lr() for imgs, labels in tqdm(iter(self.loader)): imgs = (imgs[:, (2, 1, 0)].to(conf.device) * 255) # RGB labels = labels.to(conf.device) self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = self.focalLoss(thetas, labels) loss.backward() running_loss += loss.item() / conf.batch_size self.optimizer.step() if self.step % self.board_loss_every == 0 and self.step != 0: loss_board = running_loss / self.board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % self.evaluate_every == 0 and self.step != 0: self.model.eval() for bmk in [ 'agedb_30', 'lfw', 'calfw', 'cfp_ff', 'cfp_fp', 'cplfw', 'vgg2_fp' ]: acc = eval_emore_bmk(conf, self.model, bmk) self.board_val(bmk, acc) self.model.train() if self.step % self.save_every == 0 and self.step != 0: self.save_state(conf, acc) self.step += 1 self.save_state(conf, acc, to_save_folder=True, extra='final') def myValidation(self, conf): self.model.eval() for bmk in [ 'agedb_30', 'lfw', 'calfw', 'cfp_ff', 'cfp_fp', 'cplfw', 'vgg2_fp' ]: eval_emore_bmk(conf, self.model, bmk) def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer)
class face_learner(object): def __init__(self, conf, inference=False, transfer=0): pprint.pprint(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: if transfer == 3: self.optimizer = optim.SGD( [{ 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) elif transfer == 2: self.optimizer = optim.SGD([ { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr, momentum=conf.momentum) elif transfer == 1: self.optimizer = optim.SGD([ { 'params': [self.head.kernel], 'weight_decay': 4e-4 }, ], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 5 # originally, 100 self.evaluate_every = len(self.loader) // 5 # originally, 10 self.save_every = len(self.loader) // 2 # originally, 5 # self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(self.loader.dataset.root.parent) self.val_112, self.val_112_issame = get_val_pair( self.loader.dataset.root.parent, 'val_112') else: self.threshold = conf.threshold def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.model.state_dict(), save_path / ('model_{}_accuracy:{:0.2f}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.head.state_dict(), save_path / ('head_{}_accuracy:{:0.2f}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{:0.2f}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) def load_state(self, conf, fixed_str, from_save_folder=False, model_only=False): if from_save_folder: save_path = conf.save_path else: save_path = conf.model_path self.model.load_state_dict( torch.load(save_path / 'model_{}'.format(fixed_str), map_location=conf.device)) if not model_only: self.head.load_state_dict( torch.load(save_path / 'head_{}'.format(fixed_str))) self.optimizer.load_state_dict( torch.load(save_path / 'optimizer_{}'.format(fixed_str))) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) # self.writer.add_scalar('{}_val:true accept ratio'.format(db_name), val, self.step) # self.writer.add_scalar('{}_val_std'.format(db_name), val_std, self.step) # self.writer.add_scalar('{}_far:False Acceptance Ratio'.format(db_name), far, self.step) def evaluate(self, conf, carray, issame, nrof_folds=5, tta=False): self.model.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size]) with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch) else: embeddings[idx:idx + conf.batch_size] = self.model( batch.to(conf.device)).cpu() idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:] = l2_norm(emb_batch) else: embeddings[idx:] = self.model(batch.to(conf.device)).cpu() tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds) buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor def find_lr(self, conf, init_value=1e-8, final_value=10., beta=0.98, bloding_scale=3., num=None): if not num: num = len(self.loader) mult = (final_value / init_value)**(1 / num) lr = init_value for params in self.optimizer.param_groups: params['lr'] = lr self.model.train() avg_loss = 0. best_loss = 0. batch_num = 0 losses = [] log_lrs = [] for i, (imgs, labels) in tqdm(enumerate(self.loader), total=num): imgs = imgs.to(conf.device) labels = labels.to(conf.device) batch_num += 1 self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) #Compute the smoothed loss avg_loss = beta * avg_loss + (1 - beta) * loss.item() self.writer.add_scalar('avg_loss', avg_loss, batch_num) smoothed_loss = avg_loss / (1 - beta**batch_num) self.writer.add_scalar('smoothed_loss', smoothed_loss, batch_num) #Stop if the loss is exploding if batch_num > 1 and smoothed_loss > bloding_scale * best_loss: print('exited with best_loss at {}'.format(best_loss)) plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses #Record the best loss if smoothed_loss < best_loss or batch_num == 1: best_loss = smoothed_loss #Store the values losses.append(smoothed_loss) log_lrs.append(math.log10(lr)) self.writer.add_scalar('log_lr', math.log10(lr), batch_num) #Do the SGD step #Update the lr for the next step loss.backward() self.optimizer.step() lr *= mult for params in self.optimizer.param_groups: params['lr'] = lr if batch_num > num: plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses def train(self, conf, epochs, ext='final'): self.model.train() running_loss = 0. for e in range(epochs): print('epoch {} started'.format(e)) if e == self.milestones[0]: self.schedule_lr() if e == self.milestones[1]: self.schedule_lr() if e == self.milestones[2]: self.schedule_lr() for imgs, labels in tqdm(iter(self.loader)): imgs = imgs.to(conf.device) labels = labels.to(conf.device) self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) loss = conf.ce_loss(thetas, labels) loss.backward() running_loss += loss.item() self.optimizer.step() if self.step % self.board_loss_every == 0 and self.step != 0: loss_board = running_loss / self.board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % self.evaluate_every == 0 and self.step != 0: # accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.agedb_30, self.agedb_30_issame) # self.board_val('agedb_30', accuracy, best_threshold, roc_curve_tensor) # accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.lfw, self.lfw_issame) # self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor) # accuracy, best_threshold, roc_curve_tensor = self.evaluate(conf, self.cfp_fp, self.cfp_fp_issame) # self.board_val('cfp_fp', accuracy, best_threshold, roc_curve_tensor) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.val_112, self.val_112_issame) self.board_val('n+n_val_112', accuracy, best_threshold, roc_curve_tensor) self.model.train() # if self.step % self.save_every == 0 and self.step != 0: # self.save_state(conf, accuracy, extra=ext) self.step += 1 # self.save_state(conf, accuracy, to_save_folder=True, extra=ext, model_only=True) def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer) def infer(self, conf, faces, target_embs, tta=False): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model( conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model( conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append( self.model( conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose( 1, 0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum def binfer(self, conf, faces, target_embs, tta=False): ''' return raw scores for every class faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model( conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model( conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append( self.model( conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose( 1, 0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) # print(dist) return dist.detach().cpu().numpy()
class face_learner(object): def __init__(self, conf, inference=False, embedding_size=512): conf.embedding_size = embedding_size print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).cuda() else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).cuda() print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) parameter_num_cal(self.model) self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.step = 0 self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data( self.loader.dataset.root.parent) self.writer = SummaryWriter(conf.log_path) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 if conf.multi_sphere: if conf.arcface_loss: self.head = ArcfaceMultiSphere( embedding_size=conf.embedding_size, classnum=self.class_num, num_shpere=conf.num_sphere, m=conf.m).to(conf.device) elif conf.am_softmax: self.head = MultiAm_softmax( embedding_size=conf.embedding_size, classnum=self.class_num, num_sphere=conf.num_sphere, m=conf.m).to(conf.device) else: self.head = MultiSphereSoftmax( embedding_size=conf.embedding_size, classnum=self.class_num, num_sphere=conf.num_sphere).to(conf.device) else: if conf.arcface_loss: self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) elif conf.am_softmax: self.head = Am_softmax(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) else: self.head = Softmax(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: if conf.multi_sphere: self.optimizer = optim.SGD([{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + self.head.kernel_list, 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: if conf.multi_sphere: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + self.head.kernel_list, 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 100 self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 5 self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data( self.loader.dataset.root.parent) else: self.threshold = conf.threshold def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.model.state_dict(), save_path / ('model_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.head.state_dict(), save_path / ('head_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) def get_new_state(self, path): state_dict = torch.load(path) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): if 'module' not in k: k = 'module.' + k else: k = k.replace('features.module.', 'module.features.') new_state_dict[k] = v return new_state_dict def load_state(self, save_path, fixed_str, model_only=False): self.model.load_state_dict( torch.load(save_path / 'model_{}'.format(fixed_str))) if not model_only: self.head.load_state_dict( torch.load(save_path / 'head_{}'.format(fixed_str))) self.optimizer.load_state_dict( torch.load(save_path / 'optimizer_{}'.format(fixed_str))) print(self.optimizer) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor, angle_info): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) self.writer.add_scalar('{}_same_pair_angle_mean'.format(db_name), angle_info['same_pair_angle_mean'], self.step) self.writer.add_scalar('{}_same_pair_angle_var'.format(db_name), angle_info['same_pair_angle_var'], self.step) self.writer.add_scalar('{}_diff_pair_angle_mean'.format(db_name), angle_info['diff_pair_angle_mean'], self.step) self.writer.add_scalar('{}_diff_pair_angle_var'.format(db_name), angle_info['diff_pair_angle_var'], self.step) def evaluate(self, conf, carray, issame, nrof_folds=10, tta=False, n=1): self.model.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size // n]) i = 0 with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:idx + conf.batch_size] = l2_norm( emb_batch).cpu()[:, i * conf.embedding_size // n:(i + 1) * conf.embedding_size // n] else: embeddings[idx:idx + conf.batch_size] = self.model( batch.to(conf.device)).cpu()[:, i * conf.embedding_size // n:(i + 1) * conf.embedding_size // n] idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.model(batch.to(conf.device)) + self.model( fliped.to(conf.device)) embeddings[idx:] = l2_norm( emb_batch).cpu()[:, i * conf.embedding_size // n:(i + 1) * conf.embedding_size // n] else: embeddings[idx:] = self.model(batch.to( conf.device)).cpu()[:, i * conf.embedding_size // n:(i + 1) * conf.embedding_size // n] tpr, fpr, accuracy, best_thresholds, angle_info = evaluate( embeddings, issame, nrof_folds) buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean( ), roc_curve_tensor, angle_info def find_lr(self, conf, init_value=1e-8, final_value=10., beta=0.98, bloding_scale=3., num=None): if not num: num = len(self.loader) mult = (final_value / init_value)**(1 / num) lr = init_value for params in self.optimizer.param_groups: params['lr'] = lr self.model.train() avg_loss = 0. best_loss = 0. batch_num = 0 losses = [] log_lrs = [] for i, (imgs, labels) in tqdm(enumerate(self.loader), total=num): imgs = imgs.to(conf.device) labels = labels.to(conf.device) batch_num += 1 self.optimizer.zero_grad() embeddings = self.model(imgs) thetas = self.head(embeddings, labels) if conf.multi_sphere: loss = conf.ce_loss(thetas[0], labels) for theta in thetas[1:]: loss = loss + conf.ce_loss(theta, labels) else: loss = conf.ce_loss(thetas, labels) #Compute the smoothed loss avg_loss = beta * avg_loss + (1 - beta) * loss.item() self.writer.add_scalar('avg_loss', avg_loss, batch_num) smoothed_loss = avg_loss / (1 - beta**batch_num) self.writer.add_scalar('smoothed_loss', smoothed_loss, batch_num) #Stop if the loss is exploding if batch_num > 1 and smoothed_loss > bloding_scale * best_loss: print('exited with best_loss at {}'.format(best_loss)) plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses #Record the best loss if smoothed_loss < best_loss or batch_num == 1: best_loss = smoothed_loss #Store the values losses.append(smoothed_loss) log_lrs.append(math.log10(lr)) self.writer.add_scalar('log_lr', math.log10(lr), batch_num) #Do the SGD step #Update the lr for the next step loss.backward() self.optimizer.step() lr *= mult for params in self.optimizer.param_groups: params['lr'] = lr if batch_num > num: plt.plot(log_lrs[10:-5], losses[10:-5]) return log_lrs, losses def model_evaluation(self, conf): self.model.load_state_dict(torch.load(conf.pretrained_model_path)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.agedb_30, self.agedb_30_issame, tta=True) print('age_db_acc:', accuracy) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.lfw, self.lfw_issame, tta=True) print('lfw_acc:', accuracy) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.cfp_fp, self.cfp_fp_issame, tta=True) print('cfp_acc:', accuracy) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.calfw, self.calfw_issame, tta=True) print('calfw_acc:', accuracy) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.cplfw, self.cplfw_issame, tta=True) print('cplfw_acc:', accuracy) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.vgg2_fp, self.vgg2_fp_issame, tta=True) print('vgg2_acc:', accuracy) def train(self, conf, epochs): self.model.train() running_loss = 0. if conf.pretrain: self.model_evaluation(conf) sys.exit(0) logging.basicConfig( filename=conf.log_path / 'log.txt', level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s", datefmt='%Y-%m-%d %H:%M:%S %a') logging.info( '\n******\nnum of sphere is: {},\nnet is: {},\ndepth is: {},\nlr is: {},\nbatch size is: {}\n******' .format(conf.num_sphere, conf.net_mode, conf.net_depth, conf.lr, conf.batch_size)) for e in range(epochs): print('epoch {} started,all is {}'.format(e, epochs)) if e == self.milestones[0]: self.schedule_lr() if e == self.milestones[1]: self.schedule_lr() if e == self.milestones[2]: self.schedule_lr() for imgs, labels in tqdm(iter(self.loader)): self.model.train() imgs = imgs.to(conf.device) labels = labels.to(conf.device) embeddings = self.model(imgs) thetas = self.head(embeddings, labels) if conf.multi_sphere: loss = conf.ce_loss(thetas[0], labels) for theta in thetas[1:]: loss = loss + conf.ce_loss(theta, labels) else: loss = conf.ce_loss(thetas, labels) running_loss += loss.item() self.optimizer.zero_grad() loss.backward() self.optimizer.step() if self.step % self.board_loss_every == 0 and self.step != 0: loss_board = running_loss / self.board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % self.evaluate_every == 0 and self.step != 0: accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.agedb_30, self.agedb_30_issame) print('age_db_acc:', accuracy) self.board_val('agedb_30', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('agedb_30 acc: {}'.format(accuracy)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.lfw, self.lfw_issame) print('lfw_acc:', accuracy) self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('lfw acc: {}'.format(accuracy)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.cfp_fp, self.cfp_fp_issame) print('cfp_acc:', accuracy) self.board_val('cfp', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('cfp acc: {}'.format(accuracy)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.calfw, self.calfw_issame) print('calfw_acc:', accuracy) self.board_val('calfw', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('calfw acc: {}'.format(accuracy)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.cplfw, self.cplfw_issame) print('cplfw_acc:', accuracy) self.board_val('cplfw', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('cplfw acc: {}'.format(accuracy)) accuracy, best_threshold, roc_curve_tensor, angle_info = self.evaluate( conf, self.vgg2_fp, self.vgg2_fp_issame) print('vgg2_acc:', accuracy) self.board_val('vgg2', accuracy, best_threshold, roc_curve_tensor, angle_info) logging.info('vgg2_fp acc: {}'.format(accuracy)) self.model.train() self.step += 1 def schedule_lr(self): for params in self.optimizer_corr.param_groups: params['lr'] /= 10 for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer) def infer(self, conf, faces, target_embs, tta=False): ''' faces : list of PIL Image target_embs : [n, 512] computed embeddings of faces in facebank names : recorded names of faces in facebank tta : test time augmentation (hfilp, that's all) ''' embs = [] for img in faces: if tta: mirror = trans.functional.hflip(img) emb = self.model( conf.test_transform(img).to(conf.device).unsqueeze(0)) emb_mirror = self.model( conf.test_transform(mirror).to(conf.device).unsqueeze(0)) embs.append(l2_norm(emb + emb_mirror)) else: embs.append( self.model( conf.test_transform(img).to(conf.device).unsqueeze(0))) source_embs = torch.cat(embs) diff = source_embs.unsqueeze(-1) - target_embs.transpose( 1, 0).unsqueeze(0) dist = torch.sum(torch.pow(diff, 2), dim=1) minimum, min_idx = torch.min(dist, dim=1) min_idx[minimum > self.threshold] = -1 # if no match, set idx to -1 return min_idx, minimum
class face_learner(object): def __init__(self, conf, inference=False): self.backbone = Backbone().to(conf.device) self.idprehead = PreheadID().to(conf.device) self.idhead = Arcface().to(conf.device) self.attrhead = Attrhead().to(conf.device) print('model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones train_dataset = CelebA( 'dataset', 'celebA_train.txt', trans.Compose([ trans.RandomHorizontalFlip(), trans.ToTensor(), trans.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])) valid_dataset = CelebA( 'dataset', 'celebA_validation.txt', trans.Compose([ trans.RandomHorizontalFlip(), trans.ToTensor(), trans.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])) self.loader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, pin_memory=conf.pin_memory, num_workers=conf.num_workers) self.valid_loader = DataLoader(valid_dataset, batch_size=conf.batch_size, shuffle=True, pin_memory=conf.pin_memory, num_workers=conf.num_workers) self.writer = SummaryWriter(conf.log_path) self.step = 0 paras_only_bn_1, paras_wo_bn_1 = separate_bn_paras(self.backbone) paras_only_bn_2, paras_wo_bn_2 = separate_bn_paras(self.idprehead) paras_only_bn_3, paras_wo_bn_3 = separate_bn_paras(self.attrhead) paras_only_bn = paras_only_bn_1 + paras_only_bn_2 + paras_only_bn_3 paras_wo_bn = paras_wo_bn_1 + paras_wo_bn_2 + paras_wo_bn_3 self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.idhead.kernel], 'weight_decay': 1e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 8 self.evaluate_every = len(self.loader) // 4 self.save_every = len(self.loader) // 2 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( Path("data/faces_emore/")) else: self.threshold = conf.threshold def save_state(self, conf, accuracy, to_save_folder=False, extra=None, model_only=False): if to_save_folder: save_path = conf.save_path else: save_path = conf.model_path torch.save( self.backbone.state_dict(), save_path / ('backbone_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.idprehead.state_dict(), save_path / ('idprehead_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) if not model_only: torch.save( self.idhead.state_dict(), save_path / ('idhead_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.attrhead.state_dict(), save_path / ('attrhead_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) torch.save( self.optimizer.state_dict(), save_path / ('optimizer_{}_accuracy:{}_step:{}_{}.pth'.format( get_time(), accuracy, self.step, extra))) # def load_state(self, conf, fixed_str, from_save_folder=False, model_only=False): # if from_save_folder: # save_path = conf.save_path # else: # save_path = conf.model_path # self.model.load_state_dict(torch.load(save_path/'model_{}'.format(fixed_str))) # if not model_only: # self.head.load_state_dict(torch.load(save_path/'head_{}'.format(fixed_str))) # self.optimizer.load_state_dict(torch.load(save_path/'optimizer_{}'.format(fixed_str))) def board_val(self, db_name, accuracy, best_threshold, roc_curve_tensor): self.writer.add_scalar('{}_accuracy'.format(db_name), accuracy, self.step) self.writer.add_scalar('{}_best_threshold'.format(db_name), best_threshold, self.step) self.writer.add_image('{}_roc_curve'.format(db_name), roc_curve_tensor, self.step) # self.writer.add_scalar('{}_val:true accept ratio'.format(db_name), val, self.step) # self.writer.add_scalar('{}_val_std'.format(db_name), val_std, self.step) # self.writer.add_scalar('{}_far:False Acceptance Ratio'.format(db_name), far, self.step) def evaluate(self, conf, carray, issame, nrof_folds=5, tta=False): self.backbone.eval() self.idprehead.eval() idx = 0 embeddings = np.zeros([len(carray), conf.embedding_size]) with torch.no_grad(): while idx + conf.batch_size <= len(carray): batch = torch.tensor(carray[idx:idx + conf.batch_size]) if tta: fliped = hflip_batch(batch) emb_batch = self.idprehead( self.backbone(batch.to(conf.device))) + self.idprehead( self.backbone(fliped.to(conf.device))) embeddings[idx:idx + conf.batch_size] = l2_norm(emb_batch) else: embeddings[idx:idx + conf.batch_size] = self.idprehead( self.backbone(batch.to(conf.device))).cpu() idx += conf.batch_size if idx < len(carray): batch = torch.tensor(carray[idx:]) if tta: fliped = hflip_batch(batch) emb_batch = self.idprehead( self.backbone(batch.to(conf.device))) + self.idprehead( self.backbone(fliped.to(conf.device))) embeddings[idx:] = l2_norm(emb_batch) else: embeddings[idx:] = self.idprehead( self.backbone(batch.to(conf.device))).cpu() tpr, fpr, accuracy, best_thresholds = evaluate(embeddings, issame, nrof_folds) buf = gen_plot(fpr, tpr) roc_curve = Image.open(buf) roc_curve_tensor = trans.ToTensor()(roc_curve) return accuracy.mean(), best_thresholds.mean(), roc_curve_tensor # def find_lr(self, # conf, # init_value=1e-8, # final_value=10., # beta=0.98, # bloding_scale=3., # num=None): # if not num: # num = len(self.loader) # mult = (final_value / init_value)**(1 / num) # lr = init_value # for params in self.optimizer.param_groups: # params['lr'] = lr # self.model.train() # avg_loss = 0. # best_loss = 0. # batch_num = 0 # losses = [] # log_lrs = [] # for i, (imgs, labels) in tqdm(enumerate(self.loader), total=num): # imgs = imgs.to(conf.device) # labels = labels.to(conf.device) # batch_num += 1 # self.optimizer.zero_grad() # embeddings = self.model(imgs) # thetas = self.head(embeddings, labels) # loss = conf.ce_loss(thetas, labels) # #Compute the smoothed loss # avg_loss = beta * avg_loss + (1 - beta) * loss.item() # self.writer.add_scalar('avg_loss', avg_loss, batch_num) # smoothed_loss = avg_loss / (1 - beta**batch_num) # self.writer.add_scalar('smoothed_loss', smoothed_loss,batch_num) # #Stop if the loss is exploding # if batch_num > 1 and smoothed_loss > bloding_scale * best_loss: # print('exited with best_loss at {}'.format(best_loss)) # plt.plot(log_lrs[10:-5], losses[10:-5]) # return log_lrs, losses # #Record the best loss # if smoothed_loss < best_loss or batch_num == 1: # best_loss = smoothed_loss # #Store the values # losses.append(smoothed_loss) # log_lrs.append(math.log10(lr)) # self.writer.add_scalar('log_lr', math.log10(lr), batch_num) # #Do the SGD step # #Update the lr for the next step # loss.backward() # self.optimizer.step() # lr *= mult # for params in self.optimizer.param_groups: # params['lr'] = lr # if batch_num > num: # plt.plot(log_lrs[10:-5], losses[10:-5]) # return log_lrs, losses def train(self, conf, epochs): self.backbone.train() self.idprehead.train() self.attrhead.train() running_loss = 0. for e in range(epochs): print('epoch {} started'.format(e)) if e == self.milestones[0]: self.schedule_lr() if e == self.milestones[1]: self.schedule_lr() if e == self.milestones[2]: self.schedule_lr() for imgs, labels in tqdm(iter(self.loader)): imgs = imgs.to(conf.device) labels = labels.to(conf.device) attributes = labels[:, :40] attributes = (attributes + 1) * 0.5 ids = labels[:, 40] self.optimizer.zero_grad() embeddings = self.backbone(imgs) thetas = self.idhead(self.idprehead(embeddings), ids) # attrs = self.attrhead(embeddings) # attributes = attributes.type_as(attrs) loss = conf.ce_loss(thetas, ids) loss.backward() running_loss += loss.item() self.optimizer.step() if self.step % self.board_loss_every == 0 and self.step != 0: loss_board = running_loss / self.board_loss_every self.writer.add_scalar('train_loss', loss_board, self.step) running_loss = 0. if self.step % self.evaluate_every == 0 and self.step != 0: accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.agedb_30, self.agedb_30_issame) self.board_val('agedb_30', accuracy, best_threshold, roc_curve_tensor) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.lfw, self.lfw_issame) self.board_val('lfw', accuracy, best_threshold, roc_curve_tensor) accuracy, best_threshold, roc_curve_tensor = self.evaluate( conf, self.cfp_fp, self.cfp_fp_issame) self.board_val('cfp_fp', accuracy, best_threshold, roc_curve_tensor) # attr_loss, attr_accu = self.validate_attr(conf) # print(attr_loss, attr_accu) self.backbone.train() self.idprehead.train() self.attrhead.train() if self.step % self.save_every == 0 and self.step != 0: self.save_state(conf, accuracy) self.step += 1 self.save_state(conf, accuracy, to_save_folder=True, extra='final') def schedule_lr(self): for params in self.optimizer.param_groups: params['lr'] /= 10 print(self.optimizer) def validate_attr(self, conf): self.backbone.eval() self.attrhead.eval() losses = [] accuracies = [] with torch.no_grad(): for i, (input, target) in enumerate(self.valid_loader): input = input.to(conf.device) target = target.to(conf.device) target = target[:, :40] target = (target + 1) * 0.5 target = target.type(torch.cuda.FloatTensor) # compute output embedding = self.backbone(input) output = self.attrhead(embedding) # measure accuracy and record loss loss = conf.bc_loss(output, target) pred = torch.where(torch.sigmoid(output) > 0.5, 1.0, 0.0) accuracies.append((pred == target).float().sum() / (target.size()[0] * target.size()[1])) losses.append(loss) loss_avg = sum(losses) / len(losses) accu_avg = sum(accuracies) / len(accuracies) return loss_avg, accu_avg
def __init__(self, conf, inference=False, embedding_size=512): conf.embedding_size = embedding_size print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).cuda() else: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).cuda() print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) parameter_num_cal(self.model) self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.step = 0 self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data( self.loader.dataset.root.parent) self.writer = SummaryWriter(conf.log_path) if not inference: self.milestones = conf.milestones self.loader, self.class_num = get_train_loader(conf) self.writer = SummaryWriter(conf.log_path) self.step = 0 if conf.multi_sphere: if conf.arcface_loss: self.head = ArcfaceMultiSphere( embedding_size=conf.embedding_size, classnum=self.class_num, num_shpere=conf.num_sphere, m=conf.m).to(conf.device) elif conf.am_softmax: self.head = MultiAm_softmax( embedding_size=conf.embedding_size, classnum=self.class_num, num_sphere=conf.num_sphere, m=conf.m).to(conf.device) else: self.head = MultiSphereSoftmax( embedding_size=conf.embedding_size, classnum=self.class_num, num_sphere=conf.num_sphere).to(conf.device) else: if conf.arcface_loss: self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) elif conf.am_softmax: self.head = Am_softmax(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) else: self.head = Softmax(embedding_size=conf.embedding_size, classnum=self.class_num).to( conf.device) paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: if conf.multi_sphere: self.optimizer = optim.SGD([{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + self.head.kernel_list, 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: if conf.multi_sphere: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + self.head.kernel_list, 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = len(self.loader) // 100 self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 5 self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data( self.loader.dataset.root.parent) else: self.threshold = conf.threshold
from model import Backbone, Arcface, MobileFaceNet, Am_softmax, l2_norm from torchvision import transforms as trans import PIL.Image as Image from mtcnn import MTCNN import torch import cv2 import os img_root_dir = '../img_align_celeba' save_path = '../celeba_64' device = torch.device('cuda:0') mtcnn = MTCNN() model = Backbone(50, 0.6, 'ir_se').to(device) model.eval() model.load_state_dict(torch.load('./saved_models/model_ir_se50.pth')) # threshold = 1.54 test_transform = trans.Compose([ trans.ToTensor(), trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ]) # decoder = libnvjpeg.py_NVJpegDecoder() ind = 0 embed_map = {} for root, dirs, files in os.walk(img_root_dir): files_len = len(files)
def __init__(self, conf, inference=False, need_loader=True): print(conf) if conf.use_mobilfacenet: # self.model = MobileFaceNet(conf.embedding_size).to(conf.device) self.model = torch.nn.DataParallel( MobileFaceNet(conf.embedding_size)).cuda() print('MobileFaceNet model generated') else: # self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) self.model = torch.nn.DataParallel( Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode)).cuda() print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if not inference: self.milestones = conf.milestones if need_loader: # self.loader, self.class_num = get_train_loader(conf) self.dataset = Dataset2() self.loader = DataLoader(self.dataset, batch_size=conf.batch_size, num_workers=conf.num_workers, shuffle=True, pin_memory=True) # self.loader = Loader2(conf) self.class_num = 85164 print(self.class_num, 'classes, load ok ') else: import copy conf_t = copy.deepcopy(conf) conf_t.data_mode = 'emore' self.loader, self.class_num = get_train_loader(conf_t) print(self.class_num) self.class_num = 85164 lz.mkdir_p(conf.log_path, delete=True) self.writer = SummaryWriter(conf.log_path) self.step = 0 if conf.loss == 'arcface': self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) elif conf.loss == 'softmax': self.head = MySoftmax(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) else: raise ValueError(f'{conf.loss}') print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn[:-1], 'weight_decay': 4e-5 }, { 'params': [paras_wo_bn[-1]] + [self.head.kernel], 'weight_decay': 4e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) else: self.optimizer = optim.SGD( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, momentum=conf.momentum) print(self.optimizer) # self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True) print('optimizers generated') self.board_loss_every = 100 # len(self.loader) // 100 self.evaluate_every = len(self.loader) // 10 self.save_every = len(self.loader) // 5 self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( self.loader.dataset.root_path) else: self.threshold = conf.threshold