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 __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
def __init__(self, conf, inference=False): print(conf) if conf.use_seesawFaceNet: self.model = seesaw_shuffleFaceNetv3(conf.embedding_size).to(conf.device) print('seesawFaceNetv3 (with slim) 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 = LiArcFace(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) print('using LiArcFace as loss function') print('two model heads generated') paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) self.model = torch.nn.DataParallel(self.model).cuda() cudnn.enabled = True cudnn.benchmark = True if conf.use_seesawFaceNet: print('setting optimizer for seesawFacenet') self.optimizer = optim.SGD([ {'params': paras_wo_bn[:-1], 'weight_decay': 4e-5}, {'params': [paras_wo_bn[-1]] + [self.head.weight], '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.weight], '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)//2 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) else: self.threshold = conf.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(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.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) else: self.threshold = conf.threshold
def train(epochs): logging.debug("Prepare Data") all_class = [i for i in os.listdir(data_folder)] class_num = len(all_class) regular_class = [] for idx, i in enumerate(all_class): if len(os.listdir(os.path.join(data_folder, i))) > 20: regular_class.append(i) if idx % 20000 == 0: print("Processing...", idx) loader, class_num = get_train_loader(data_folder_stage2, batch_size) agedb_30, cfp_fp, lfw, agedb_30_issame, cfp_fp_issame, lfw_issame = get_val_data( val_folder) ## Load Model enc = Encoder().to(device) dec = Decoder().to(device) r = Distillation_R().to(device) head = FC_softmax(320, class_num).to(device) enc.load_state_dict( torch.load( os.path.join(pretrained_model_path, 'enc_{}_{}.pth'.format(assigned_epoch, assigned_step)))) dec.load_state_dict( torch.load( os.path.join(pretrained_model_path, 'dec_{}_{}.pth'.format(assigned_epoch, assigned_step)))) r.load_state_dict( torch.load( os.path.join(pretrained_model_path, 'r_{}_{}.pth'.format(assigned_epoch, assigned_step)))) head.load_state_dict( torch.load( os.path.join( pretrained_model_path, 'head_{}_{}.pth'.format(assigned_epoch, assigned_step)))) enc.eval() dec.eval() r.eval() head.eval() model = [enc, dec, r] ## Set Training Criterion ce_loss = nn.CrossEntropyLoss() l2_loss = nn.MSELoss() optimizer = optim.Adam(model[0].parameters(), lr=0.00001) optimizer.add_param_group({'params': model[1].parameters()}) optimizer.add_param_group({'params': model[2].parameters()}) optimizer.add_param_group({'params': head.parameters()}) logging.debug("Start Training") ## Initial Training running_loss = 0 step = 0 acc_max = 0 for e in range(epochs): for stage in [1, 2]: ## Initial Train For Stage 1 and 2 if stage == 1: center, Q, h, ur_class_file = None, None, None, None center, Q, h, ur_class_file = UpdateStats( model[0], data_folder, all_class, regular_class) logging.debug( "Center_num:{}, h_num:{}, ur_class_num:{}".format( len(center), len(h), len(ur_class_file))) else: iterr = iter(loader) ## Start Training for step_stage in range(20000): if stage == 1: regular_batch = random.sample(range(len(h)), batch_size) ur_batch = random.sample(range(len(ur_class_file)), batch_size) g_r_list, g_u_list, g_t_list, label_r_list, label_u_list, label_t_list = list( ), list(), list(), list(), list(), list() for idx_r, idx_u in zip(regular_batch, ur_batch): ##Prepare data for First Batch () label_r, g_r = h[idx_r] c_r = center[label_r] g_r_list.append(g_r) label_r_list.append(torch.tensor(int(label_r))) ##Prepare data for Second Batch label_u, g_u = ur_class_file[idx_u] c_u = center[label_u] g_u_list.append(g_u) label_u_list.append(torch.tensor(int(label_u))) ##Prepare data for Third Batch g_t = c_u + torch.mm(Q, (g_r - c_r).t()).t() g_t_list.append(g_t.view(-1)) label_t_list.append(torch.tensor(int(label_u))) g_r_list = torch.stack(g_r_list).to(device) label_r_list = torch.stack(label_r_list).to(device) g_u_list = torch.stack(g_u_list).to(device) label_u_list = torch.stack(label_u_list).to(device) g_t_list = torch.stack(g_t_list).to(device) label_t_list = torch.stack(label_t_list).to(device) for g, labels in zip( [g_r_list, g_u_list, g_t_list], [label_r_list, label_u_list, label_t_list]): optimizer.zero_grad() embs = model[2](g) fc_out = head(embs) loss_ce = ce_loss(fc_out, labels) loss_reg = l2_norm(fc_out) loss = loss_ce + loss_reg * 0.25 loss.backward() running_loss += loss.item() optimizer.step() running_loss /= 3 elif stage == 2: imgs, labels = next(iterr) imgs = imgs.to(device) labels = labels.to(device) optimizer.zero_grad() g = model[0](imgs) img_decs = model[1](g) with torch.no_grad(): embs = model[2](g) fc_out = head(embs) loss_ce = ce_loss(fc_out, labels) loss_mse = l2_loss(imgs, img_decs) loss_reg = l2_norm(fc_out) loss = loss_ce + loss_mse + loss_reg * 0.25 loss.backward() running_loss += loss.item() optimizer.step() if step % board_loss_every == 0 and step != 0: loss_board = running_loss / board_loss_every printout = "stage:{}, step:{}, epoch:{}, train_loss:{}".format( stage, step, e, loss_board) logging.debug(printout) running_loss = 0 if step % evaluate_every == 0 and step != 0: accuracy, best_threshold = verify(model, agedb_30, agedb_30_issame) printout = "dataset:age30db, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) accuracy, best_threshold = verify(model, lfw, lfw_issame) printout = "dataset:lfw, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) accuracy, best_threshold = verify(model, cfp_fp, cfp_fp_issame) printout = "dataset:cfp_fp, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) if accuracy > acc_max: torch.save( model[0].state_dict(), '{}/enc_alter_{}_{}.pth'.format( save_path, e, step)) torch.save( model[1].state_dict(), '{}/dec_alter_{}_{}.pth'.format( save_path, e, step)) torch.save( model[2].state_dict(), '{}/r_alter_{}_{}.pth'.format(save_path, e, step)) torch.save( head.state_dict(), '{}/head_alter_{}_{}.pth'.format( save_path, e, step)) acc_max = accuracy logging.debug("Save ckpt at epoch:{} step:{}".format( e, step)) step += 1
def __init__(self, conf, train=True): make_dir(conf.work_path) make_dir(conf.model_path) make_dir(conf.log_path) if conf.gpu_ids: assert torch.cuda.is_available(), 'GPU is not avalialble!' torch.backends.cudnn.benckmark = True conf.device = torch.device('cuda') else: conf.device = torch.device('cpu') self.gpu_ids = conf.gpu_ids self.model = None self.net_type = '{}_{}'.format(conf.net_mode, conf.net_depth) if conf.net_mode == 'ir' or conf.net_mode == 'ir_se': self.model = IRNet(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) elif conf.net_mode == 'resnet': if conf.net_depth == 18: self.model = ResNet_18().to(conf.device) elif conf.net_depth == 34: self.model = ResNet_34().to(conf.device) elif conf.net_depth == 50: self.model = ResNet_50().to(conf.device) elif conf.net_depth == 101: self.model = ResNet_101().to(conf.device) elif conf.net_depth == 152: self.model = ResNet_152().to(conf.device) else: raise NotImplementedError( "Model {}_{} is not implemented".format( conf.net_mode, conf.net_depth)) elif conf.net_mode == 'lightcnn': if conf.net_depth == 9: self.model = LightCNN_9Layers(drop_ratio=conf.drop_ratio).to( conf.device) elif conf.net_depth == 29: self.model = LightCNN_29Layers(drop_ratio=conf.drop_ratio).to( conf.device) else: raise NotImplementedError( "Model {}_{} is not implemented".format( conf.net_mode, conf.net_depth)) else: NotImplementedError("Model {}_{} is not implemented".format( conf.net_mode, conf.net_depth)) assert self.model is not None, "Model is NONE!!" if train: 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_ADAM: self.optimizer = optim.Adam( [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }], lr=conf.lr, betas=(0.9, 0.999)) 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) # if conf.cosine_lr == True: # self.scheduler = CosineAnnealingWarmUpRestarts(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.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( conf.data_folder)
def train(epochs): logging.debug("Prepare Data") loader, class_num = get_train_loader(data_folder, batch_size) agedb_30, cfp_fp, lfw, agedb_30_issame, cfp_fp_issame, lfw_issame = get_val_data( val_folder) ## Load Model enc = Encoder().to(device) dec = Decoder().to(device) r = Distillation_R().to(device) model = [enc, dec, r] head = FC_softmax(320, class_num).to(device) if pretrained: enc.load_state_dict( torch.load( os.path.join( pretrained_model_path, 'enc_{}_{}.pth'.format(assigned_epoch, assigned_step)))) dec.load_state_dict( torch.load( os.path.join( pretrained_model_path, 'dec_{}_{}.pth'.format(assigned_epoch, assigned_step)))) r.load_state_dict( torch.load( os.path.join( pretrained_model_path, 'r_{}_{}.pth'.format(assigned_epoch, assigned_step)))) head.load_state_dict( torch.load( os.path.join( pretrained_model_path, 'head_{}_{}.pth'.format(assigned_epoch, assigned_step)))) enc.eval() dec.eval() r.eval() head.eval() model = [enc, dec, r] ## Set Training Criterion ce_loss = nn.CrossEntropyLoss() l2_loss = nn.MSELoss() optimizer = optim.Adam(model[0].parameters(), lr=0.0002) optimizer.add_param_group({'params': model[1].parameters()}) optimizer.add_param_group({'params': model[2].parameters()}) optimizer.add_param_group({'params': head.parameters()}) ## Initial Training running_loss = 0 step = 0 acc_max = 0 logging.debug("Start Training") for e in range(epochs): for imgs, labels in iter(loader): imgs = imgs.to(device) labels = labels.to(device) optimizer.zero_grad() g = model[0](imgs) img_decs = model[1](g) embs = model[2](g) fc_out = head(embs) loss_ce = ce_loss(fc_out, labels) loss_mse = l2_loss(imgs, img_decs) loss_reg = l2_norm(fc_out) loss = loss_ce + loss_mse + loss_reg * 0.25 loss.backward() running_loss += loss.item() optimizer.step() if step % board_loss_every == 0 and step != 0: loss_board = running_loss / board_loss_every printout = "step:{}, epoch:{}, train_loss:{}".format( step, e, loss_board) logging.debug(printout) running_loss = 0. if step % evaluate_every == 0 and step != 0: accuracy, best_threshold = verify(model, agedb_30, agedb_30_issame) printout = "dataset:age30db, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) accuracy, best_threshold = verify(model, lfw, lfw_issame) printout = "dataset:lfw, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) accuracy, best_threshold = verify(model, cfp_fp, cfp_fp_issame) printout = "dataset:cfp_fp, acc:{}, best_threshold:{}".format( accuracy, best_threshold) logging.debug(printout) if accuracy > acc_max: torch.save(model[0].state_dict(), 'ckpt/enc_{}_{}.pth'.format(e, step)) torch.save(model[1].state_dict(), 'ckpt/dec_{}_{}.pth'.format(e, step)) torch.save(model[2].state_dict(), 'ckpt/r_{}_{}.pth'.format(e, step)) torch.save(head.state_dict(), 'ckpt/head_{}_{}.pth'.format(e, step)) acc_max = accuracy logging.debug("Save ckpt at epoch:{} step:{}".format( e, step)) step += 1
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 self.writer = SummaryWriter(conf.log_path) self.step = 0 logger.info('loading data...') self.loader_arc, self.class_num_arc = get_train_loader( conf, 'emore', sample_identity=False) # self.loader_tri, self.class_num_tri = get_train_loader(conf, 'glint', sample_identity=True) emore_root = conf.data_path / 'train' / 'faces_emore_16_per_peron' / 'imgs' kc_root = conf.data_path / 'test' / 'kc_employee_dynamic_112' self.loader_tri, self.class_num_tri = get_train_loader_concat( conf, [emore_root, kc_root], sample_identity=True) self.head_arc = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num_arc).cuda() self.head_tri = Triplet().cuda() logger.debug('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_arc.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_arc.kernel], '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.resume: logger.debug('resume...') self.load_state(conf, 'last.pth', from_save_folder=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], find_unused_parameters=True).cuda( ) #add line for distributed self.board_loss_every = len(self.loader_arc) // 100 self.evaluate_every = len(self.loader_arc) // 2 self.save_every = len(self.loader_arc) // 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_arc.dataset.root).parent) self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data( conf.emore_folder) else: self.threshold = conf.threshold self.loader, self.query_ds, self.gallery_ds = get_test_loader(conf)
def __init__(self, conf, inference=False, train_transforms=None, val_transforms=None, train_loader=None): print(conf) if conf.use_mobilfacenet: self.model = MobileFaceNet(conf.embedding_size).to(conf.device) print('MobileFaceNet model generated') else: self.milestones = conf.milestones if train_loader is None: self.loader, self.class_num = get_train_loader( conf, train_transforms) else: self.loader = train_loader self.class_num = conf.num_classes if conf.net_mode in ['ir', 'ir_se']: self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode, conf.use_gap).to(conf.device) else: import json self.model = MetricNet(model_name=conf.net_mode, pooling=conf.pooling, use_fc=True, fc_dim=conf.embedding_size, dropout=conf.last_fc_dropout, pretrained=conf.pretrained, class_num=self.class_num).to( conf.device) print('{}_{} model generated'.format(conf.net_mode, conf.net_depth)) if conf.use_mobilfacenet or conf.net_mode in ['ir', 'ir_se']: self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num).to(conf.device) else: if conf.loss_module == 'arcface': self.head = ArcMarginProduct(self.model.final_in_features, self.class_num, s=conf.s, m=conf.margin, easy_margin=False, ls_eps=conf.ls_eps).to( conf.device) elif conf.loss_module == 'cosface': self.head = AddMarginProduct(self.model.final_in_features, self.class_num, s=conf.s, m=conf.margin).to(conf.device) elif conf.loss_module == 'adacos': self.head = AdaCos(self.model.final_in_features, self.class_num, m=conf.margin, theta_zero=conf.theta_zero).to( conf.device) else: self.head = nn.Linear(self.model.final_in_features, self.class_num).to(conf.device) print('two model heads generated') if conf.ft_model_path: self.load_ft_model(conf.ft_model_path, not conf.no_strict) elif conf.restore_suffix: self.load_state(conf, conf.restore_suffix, from_save_folder=False, model_only=False) if not inference: self.writer = SummaryWriter(conf.log_path) self.step = 0 paras_only_bn, paras_wo_bn = separate_bn_paras(self.model) if conf.use_mobilfacenet: params = [{ '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 }] wd = 4e-5 else: # if conf.net_mode in ['ir', 'ir_se']: # params = [ # {'params': paras_wo_bn + [self.head.weight], 'weight_decay': 5e-4}, # {'params': paras_only_bn} # ] params = [{ 'params': paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4 }, { 'params': paras_only_bn }] wd = 5e-4 # else: # params = self.model.parameters() # wd = conf.wd # # params = [ # # {'params': paras_wo_bn + [self.head.weight], 'weight_decay': conf.wd}, # 5e-4}, # # {'params': paras_only_bn} # # ] if conf.optimizer == 'sgd': self.optimizer = optim.SGD( params, lr=conf.lr, momentum=conf.momentum) # , weight_decay=wd) elif conf.optimizer == 'adam': self.optimizer = optim.Adam( params, lr=conf.lr) # , weight_decay=wd) 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.board_loss_every = 20 self.evaluate_every = len(self.loader) self.save_every = len(self.loader) if conf.data_mode == 'common': import json val_img_dir_map = json.loads(conf.val_img_dirs) self.val_dataloaders = {} for val_name in val_img_dir_map: val_img_dir = val_img_dir_map[val_name] val_dataloader, common_val_issame = get_common_val_data( val_img_dir, conf.max_positive_cnt, conf.val_batch_size, conf.val_pin_memory, conf.num_workers, val_transforms=val_transforms, use_pos=not conf.not_use_pos, use_neg=not conf.not_use_neg, val_smapling_type=conf.val_smapling_type, use_keras_model=conf.use_val_left_right_check) self.val_dataloaders[val_name] = [ val_dataloader, common_val_issame ] elif conf.data_mode == 'dacon_landmark': pass else: 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) else: self.threshold = conf.threshold