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
