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
