def train_sst(conf):
if conf.model_type == 'attention':
probe_net = AttentionNet(attention_stages=conf.attention_stages, dim=conf.feat_dim)
gallery_net = AttentionNet(attention_stages=conf.attention_stages, dim=conf.feat_dim)
elif conf.model_type == 'mobilefacenet':
probe_net = MobileFaceNet(conf.feat_dim)
gallery_net = MobileFaceNet(conf.feat_dim)
moving_average(probe_net, gallery_net, 0)
prototype = Prototype(conf.feat_dim, conf.queue_size, conf.scale,conf.margin, conf.loss_type).cuda()
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4)
lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.lr_decay_epochs, gamma=0.1)
probe_net = torch.nn.DataParallel(probe_net).cuda()
gallery_net = torch.nn.DataParallel(gallery_net).cuda()
train_dict = trainlist_to_dict(conf.source_file)
for epoch in range(1, conf.epochs + 1):
if epoch == 1:
curr_train_list, curr_id_list = train_sample(train_dict, conf.class_num, conf.queue_size)
else:
curr_train_list, curr_id_list = train_sample(train_dict, conf.class_num, conf.queue_size, curr_id_list)
data_loader = DataLoader(lmdb_utils.SingleLMDBDataset(conf.source_lmdb, curr_train_list, conf.key),
conf.batch_size, shuffle = False, num_workers=4, drop_last = True)
train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer,
criterion, epoch, conf)
lr_schedule.step()
def _resume_train(
self, model_path, lr
): # only for arcface, change 'metric' in config to arcface to make _learner() work
checkpoint = torch.load(model_path)
model = MobileFaceNet(self.config.embedding_size).to(self.device)
metric = ArcFace(self.embedding_size, self.class_num).to(self.device)
model = nn.DataParallel(model)
metric = nn.DataParallel(metric)
model.load_state_dict(checkpoint['net'])
net_params = add_weight_decay(model, self.config.weight_decay)
metric_params = add_weight_decay(metric, self.config.class_wd)
parameters = net_params + metric_params
optimizer = optim.SGD(parameters,
lr=lr,
momentum=self.config.momentum,
weight_decay=self.config.weight_decay)
current_lr = optimizer.param_groups[0]['lr']
self.logger.info(
'Resume train {}, metric state is arcface, basic learning rate is {:f}'
.format(model_path, current_lr))
accuracy, threshold = test.test(
self.config, model
) #f"./{self.config.checkpoints_path}/epoch{e}_{batch_idx}.pth"
self.logger.info(
'Loading model from \'{}\',\t test accuracy = {:.4f}, threshold = {:.4f}\n'
.format(model_path, accuracy, threshold))
self._learner(optimizer, model, metric)
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
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):
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 __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.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 main():
#video load
VideoPath = "../videos/3.mp4"
imagelist = Video2list(VideoPath)
#face detect
mtcnn_path = os.path.join(os.path.dirname(__file__), 'mtcnn-model')
detector = MtcnnDetector(model_folder=mtcnn_path,
ctx=mx.cpu(0),
num_worker=1,
minsize=80,
accurate_landmark=True,
threshold=[0.6, 0.7, 0.9])
Videoimginfo = []
for img in imagelist:
ret = detector.detect_face(img)
Videoimginfo.append(ret)
#face feature get
model = MobileFaceNet(512)
model_static_cnn = torch.load("model_mobilefacenet.pth",
map_location=lambda storage, loc: storage)
net_model_static_cnn = {}
for k, v in model_static_cnn.items():
if k == "fc2.weight":
continue
if k == "fc2.bias":
continue
net_model_static_cnn[k] = v
model.load_state_dict(net_model_static_cnn)
model.eval()
imageinfo = []
allFaceFeture = []
for item in range(len(imagelist)):
if Videoimginfo[item] is not None:
image = imagelist[item]
ret = Videoimginfo[item]
facefeature = Facefeature(ret, image, model)
imageinfo.append(len(facefeature[0]))
allFaceFeture += facefeature[0]
Videoimginfo[item] = [facefeature[1], facefeature[2]]
else:
imageinfo.append(0)
Facecalsslist, classnum = dbscan(np.array(allFaceFeture), distance, minPt)
print(Facecalsslist, classnum)
#pic2video
fourcc = cv2.VideoWriter_fourcc(*"MJPG")
videoWrite = cv2.VideoWriter(
'output.avi', fourcc, 25,
(imagelist[0].shape[1], imagelist[0].shape[0]))
font = cv2.FONT_HERSHEY_SIMPLEX
cc = 0
flag = 0
for item in range(len(imageinfo)):
img = imagelist[item]
if imageinfo[item] == 0:
videoWrite.write(img)
cv2.imwrite("./ll/%d.jpg" % cc, img)
else:
#in this one pic may be has more than one pic
# rectangle point lable ;
bbox, point = Videoimginfo[item]
for q in range(len(point)):
for i in range(5):
cv2.circle(img, (int(point[q][i]), (int(point[q][i + 5]))),
3, (0, 255, 0), -1)
cv2.rectangle(img, (int(bbox[q][0]), int(bbox[q][1])),
(int(bbox[q][2]), int(bbox[q][3])),
(0, 255, 255), 2)
cv2.putText(img, "%d" % Facecalsslist[flag],
(int(bbox[q][0]), int(bbox[q][1])), font, 1.2,
(255, 255, 255), 2)
flag += 1
cv2.imwrite("./ll/%d.jpg" % cc, img)
videoWrite.write(img)
cc += 1
videoWrite.release()
def __init__(self, conf, inference=False, need_loader=True):
print(conf)
if conf.use_mobilfacenet:
# self.model = MobileFaceNet(conf.embedding_size).to(conf.device)
self.model = torch.nn.DataParallel(
MobileFaceNet(conf.embedding_size)).cuda()
print('MobileFaceNet model generated')
else:
# self.model = Backbone(conf.net_depth, conf.drop_ratio, conf.net_mode).to(conf.device)
self.model = torch.nn.DataParallel(
Backbone(conf.net_depth, conf.drop_ratio,
conf.net_mode)).cuda()
print('{}_{} model generated'.format(conf.net_mode,
conf.net_depth))
if not inference:
self.milestones = conf.milestones
if need_loader:
# self.loader, self.class_num = get_train_loader(conf)
self.dataset = Dataset2()
self.loader = DataLoader(self.dataset,
batch_size=conf.batch_size,
num_workers=conf.num_workers,
shuffle=True,
pin_memory=True)
# self.loader = Loader2(conf)
self.class_num = 85164
print(self.class_num, 'classes, load ok ')
else:
import copy
conf_t = copy.deepcopy(conf)
conf_t.data_mode = 'emore'
self.loader, self.class_num = get_train_loader(conf_t)
print(self.class_num)
self.class_num = 85164
lz.mkdir_p(conf.log_path, delete=True)
self.writer = SummaryWriter(conf.log_path)
self.step = 0
if conf.loss == 'arcface':
self.head = Arcface(embedding_size=conf.embedding_size,
classnum=self.class_num).to(conf.device)
elif conf.loss == 'softmax':
self.head = MySoftmax(embedding_size=conf.embedding_size,
classnum=self.class_num).to(conf.device)
else:
raise ValueError(f'{conf.loss}')
print('two model heads generated')
paras_only_bn, paras_wo_bn = separate_bn_paras(self.model)
if conf.use_mobilfacenet:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn[:-1],
'weight_decay': 4e-5
}, {
'params': [paras_wo_bn[-1]] + [self.head.kernel],
'weight_decay': 4e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
else:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn + [self.head.kernel],
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
print(self.optimizer)
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
print('optimizers generated')
self.board_loss_every = 100 # len(self.loader) // 100
self.evaluate_every = len(self.loader) // 10
self.save_every = len(self.loader) // 5
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
self.loader.dataset.root_path)
else:
self.threshold = conf.threshold
def __init__(self, conf, inference=False, embedding_size=512):
conf.embedding_size = embedding_size
print(conf)
if conf.use_mobilfacenet:
self.model = MobileFaceNet(conf.embedding_size).cuda()
else:
self.model = Backbone(conf.net_depth, conf.drop_ratio,
conf.net_mode).cuda()
print('{}_{} model generated'.format(conf.net_mode,
conf.net_depth))
parameter_num_cal(self.model)
self.milestones = conf.milestones
self.loader, self.class_num = get_train_loader(conf)
self.step = 0
self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data(
self.loader.dataset.root.parent)
self.writer = SummaryWriter(conf.log_path)
if not inference:
self.milestones = conf.milestones
self.loader, self.class_num = get_train_loader(conf)
self.writer = SummaryWriter(conf.log_path)
self.step = 0
if conf.multi_sphere:
if conf.arcface_loss:
self.head = ArcfaceMultiSphere(
embedding_size=conf.embedding_size,
classnum=self.class_num,
num_shpere=conf.num_sphere,
m=conf.m).to(conf.device)
elif conf.am_softmax:
self.head = MultiAm_softmax(
embedding_size=conf.embedding_size,
classnum=self.class_num,
num_sphere=conf.num_sphere,
m=conf.m).to(conf.device)
else:
self.head = MultiSphereSoftmax(
embedding_size=conf.embedding_size,
classnum=self.class_num,
num_sphere=conf.num_sphere).to(conf.device)
else:
if conf.arcface_loss:
self.head = Arcface(embedding_size=conf.embedding_size,
classnum=self.class_num).to(
conf.device)
elif conf.am_softmax:
self.head = Am_softmax(embedding_size=conf.embedding_size,
classnum=self.class_num).to(
conf.device)
else:
self.head = Softmax(embedding_size=conf.embedding_size,
classnum=self.class_num).to(
conf.device)
paras_only_bn, paras_wo_bn = separate_bn_paras(self.model)
if conf.use_mobilfacenet:
if conf.multi_sphere:
self.optimizer = optim.SGD([{
'params': paras_wo_bn[:-1],
'weight_decay': 4e-5
}, {
'params': [paras_wo_bn[-1]] + self.head.kernel_list,
'weight_decay':
4e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
else:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn[:-1],
'weight_decay': 4e-5
}, {
'params': [paras_wo_bn[-1]] + [self.head.kernel],
'weight_decay': 4e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
else:
if conf.multi_sphere:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn + self.head.kernel_list,
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
else:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn + [self.head.kernel],
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
print(self.optimizer)
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, patience=40, verbose=True)
print('optimizers generated')
self.board_loss_every = len(self.loader) // 100
self.evaluate_every = len(self.loader) // 10
self.save_every = len(self.loader) // 5
self.agedb_30, self.cfp_fp, self.lfw, self.calfw, self.cplfw, self.vgg2_fp, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame, self.calfw_issame, self.cplfw_issame, self.vgg2_fp_issame = get_val_data(
self.loader.dataset.root.parent)
else:
self.threshold = conf.threshold
def __init__(self, ConfigTable):
""" Train on pytorch:
Step1: Create a dataloader to prepare train data.
Step2: Create a logger to write train imfomartion to your .log.
Step3: Choose your network backbone, metric function, loss function.
Step4: In pytorch the weight_decay will be applied to bias, batchnorm layer, change weight_decay to 0 at these params.
Step5: Define the optimizer eg. SGD, Adam
Step6: Now, let's start your train
"""
self.config = ConfigTable # Config table for train including lr_step, checkpoints
os.makedirs(self.config.checkpoints_path, exist_ok=True)
nowtime = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())
self.logger = get_logger("./" + nowtime + "-train.log")
self.logger.info(
'Successfully create train log file at \'./{}-train.log\''.format(
nowtime))
# dataloader
self.dataloader, self.class_num = load_data(self.config, training=True)
self.embedding_size = self.config.embedding_size
self.device = self.config.device
self.logger.info(
'Successfully create DataLoader. In data the total of class numbers is {} and embedding size is {}'
.format(self.class_num, self.embedding_size))
# generate network
if self.config.backbone == 'myfmobile':
self.net = MobileFaceNet(self.embedding_size).to(
self.device
) # Create net and copy net tensor to the GPU, do it before loading data
self.logger.info("Network backbone is {}".format(
self.config.backbone))
self.logger.info("{}".format(self.net))
elif self.config.backbone == 'resnet20':
self.net = ResNet20().to(self.device)
self.logger.info("Network backbone is {}".format(
self.config.backbone))
self.logger.info("{}".format(self.net))
if self.config.metric == 'arcface':
self.metric = ArcFace(self.embedding_size,
self.class_num).to(self.device)
self.logger.info("Metric fucntion is {}".format(
self.config.metric))
self.logger.info("{}".format(self.metric))
elif self.config.metric == 'softmax':
self.metric = DenseClassifier(self.embedding_size,
self.class_num).to(self.device)
self.logger.info("Metric fucntion is {}".format(
self.config.metric))
self.logger.info("{}".format(self.metric))
elif self.config.metric == 'normlinear':
self.metric = NormLinear(512, self.class_num).to(self.device)
self.logger.info("Metric fucntion is {}".format(
self.config.metric))
self.logger.info("{}".format(self.metric))
else:
self.logger.info("Please specify a metric")
exit(0)
self._weight_init()
# Send data to multiple gpu
self.net = nn.DataParallel(self.net)
self.metric = nn.DataParallel(self.metric)
# Remove weight_decay in batchnorm and convolution bias, refer to https://arxiv.org/abs/1706.05350
net_params = add_weight_decay(self.net, self.config.weight_decay)
metric_params = add_weight_decay(self.metric, self.config.class_wd)
self.parameters = net_params + metric_params
if self.config.loss == 'focal_loss':
self.criterion = FocalLoss(gamma=2)
self.logger.info("Loss function is FocalLoss")
else:
self.criterion = nn.CrossEntropyLoss()
self.logger.info("Loss function is CrossEntropyLoss")
if self.config.optimizer == 'sgd':
self.optimizer = optim.SGD(self.parameters,
lr=self.config.lr,
momentum=self.config.momentum,
weight_decay=self.config.weight_decay)
self.logger.info("Optimaizer is SGD")
self.logger.info("{}".format(self.optimizer))
else:
self.optimizer = optim.Adam(self.parameters,
lr=self.config.lr,
weight_decay=self.config.weight_decay)
self.logger.info("Optimaizer is Adam")
self.logger.info("{}".format(self.optimizer))
def __init__(self, conf, inference=False):
print(conf)
self.num_splits = int(conf.meta_file.split('_labels.txt')[0][-1])
if conf.use_mobilfacenet:
self.model = MobileFaceNet(conf.embedding_size)
print('MobileFaceNet model generated')
else:
self.model = Backbone(conf.net_depth, conf.drop_ratio,
conf.net_mode)
print('{}_{} model generated'.format(conf.net_mode,
conf.net_depth))
if conf.device > 1:
gpu_ids = list(
range(0, min(torch.cuda.device_count(), conf.device)))
self.model = nn.DataParallel(self.model, device_ids=gpu_ids).cuda()
else:
self.model = self.model.cuda()
if not inference:
self.milestones = conf.milestones
if conf.remove_single is True:
conf.meta_file = conf.meta_file.replace('.txt', '_clean.txt')
meta_file = open(conf.meta_file, 'r')
meta = meta_file.readlines()
pseudo_all = [int(item.split('\n')[0]) for item in meta]
pseudo_classnum = set(pseudo_all)
if -1 in pseudo_classnum:
pseudo_classnum = len(pseudo_classnum) - 1
else:
pseudo_classnum = len(pseudo_classnum)
print('classnum:{}'.format(pseudo_classnum))
pseudo_classes = [
pseudo_all[count[index]:count[index + 1]]
for index in range(self.num_splits)
]
meta_file.close()
train_dataset = [get_train_dataset(conf.emore_folder)] + [
get_pseudo_dataset([conf.pseudo_folder, index + 1],
pseudo_classes[index], conf.remove_single)
for index in range(self.num_splits)
]
self.class_num = [num for _, num in train_dataset]
print('Loading dataset done')
train_longest_size = [len(item[0]) for item in train_dataset]
temp = int(np.floor(conf.batch_size // (self.num_splits + 1)))
self.batch_size = [conf.batch_size - temp * self.num_splits
] + [temp] * self.num_splits
train_longest_size = max([
int(np.floor(td / bs))
for td, bs in zip(train_longest_size, self.batch_size)
])
train_sampler = [
GivenSizeSampler(td[0],
total_size=train_longest_size * bs,
rand_seed=None)
for td, bs in zip(train_dataset, self.batch_size)
]
self.train_loader = [
DataLoader(train_dataset[k][0],
batch_size=self.batch_size[k],
shuffle=False,
pin_memory=conf.pin_memory,
num_workers=conf.num_workers,
sampler=train_sampler[k])
for k in range(1 + self.num_splits)
]
print('Loading loader done')
self.writer = SummaryWriter(conf.log_path)
self.step = 0
self.head = [
Arcface(embedding_size=conf.embedding_size,
classnum=self.class_num[0]),
Arcface(embedding_size=conf.embedding_size,
classnum=pseudo_classnum)
]
if conf.device > 1:
self.head = [
nn.DataParallel(self.head[0], device_ids=gpu_ids).cuda(),
nn.DataParallel(self.head[1], device_ids=gpu_ids).cuda()
]
else:
self.head = [self.head[0].cuda(), self.head[1].cuda()]
print('two model heads generated')
paras_only_bn, paras_wo_bn = separate_bn_paras(self.model.module)
if conf.use_mobilfacenet:
self.optimizer = optim.SGD(
[{
'params': paras_wo_bn[:-1],
'weight_decay': 4e-5
}, {
'params': [paras_wo_bn[-1]] + [self.head.parameters()],
'weight_decay': 4e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
else:
params = [a.module.parameters() for a in self.head]
params = list(params[0]) + list(params[1])
#from IPython import embed;embed()
self.optimizer = optim.SGD([{
'params': paras_wo_bn + params,
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
print(self.optimizer)
if conf.resume is not None:
self.start_epoch = self.load_state(conf.resume)
else:
self.start_epoch = 0
print('optimizers generated')
self.board_loss_every = len(self.train_loader[0]) // 10
self.evaluate_every = len(self.train_loader[0]) // 5
self.save_every = len(self.train_loader[0]) // 5
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
conf.eval_path)
else:
self.threshold = conf.threshold
def __init__(self, conf, inference=False):
print(conf)
if conf.use_mobilfacenet:
self.model = MobileFaceNet(conf.embedding_size).to(conf.device)
print('MobileFaceNet model generated')
else:
self.model = Backbone(conf.net_depth, conf.drop_ratio,
conf.net_mode).to(conf.device)
print('{}_{} model generated'.format(conf.net_mode,
conf.net_depth))
if not inference:
self.milestones = conf.milestones
print('prepare train loader..')
self.loader, self.class_num = get_train_loader(conf)
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
self.writer = SummaryWriter(str(conf.log_path / current_time))
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.lrscheduler = optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, factor=0.1, patience=20, verbose=True)
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
print('optimizers generated')
# len(self.loader): number of batches?
self.board_loss_every = len(self.loader) // 120
self.evaluate_every = len(self.loader) // 40
self.save_every = len(self.loader) // 40
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
self.loader.dataset.root.parent)
else:
self.threshold = conf.threshold
pairs = lfw.read_pairs(os.path.expanduser(lfw_pairs))
# Get the paths for the corresponding images
paths, actual_issame = lfw.get_paths(os.path.expanduser(lfw_dir), pairs)
# print(paths)
print('paths len:', len(paths))
print('pairs len: ', len(actual_issame))
embeddings = np.zeros([len(paths), 512])
# arcface = ArcFace(classnum=10).to('cuda')
# backbone = arcface.backbone
# backbone.load_state_dict(torch.load('resnet50.pth'))
mobileFacenet = MobileFaceNet(512).to('cuda')
mobileFacenet.load_state_dict(torch.load('mobilefacenet.pth'))
mobileFacenet.eval()
with torch.no_grad():
for idx, path in enumerate(paths):
print('process image no:', idx)
img = Image.open(path)
embedding = get_embedding(mobileFacenet,
img,
tta=False,
device='cuda')
embeddings[idx] = embedding.cpu().numpy()
# np.save('temp2.npy', embeddings)
# embeddings = np.load('temp.npy')
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 = []
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)
self.growup = GrowUP().to(conf.device)
self.discriminator = Discriminator().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)
if conf.discriminator:
self.child_loader, self.adult_loader = get_train_loader_d(conf)
os.makedirs(conf.log_path, exist_ok=True)
self.writer = SummaryWriter(conf.log_path)
self.step = 0
self.head = Arcface(embedding_size=conf.embedding_size,
classnum=self.class_num).to(conf.device)
# Will not use anymore
if conf.use_dp:
self.model = nn.DataParallel(self.model)
self.head = nn.DataParallel(self.head)
print(self.class_num)
print(conf)
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)
if conf.discriminator:
self.optimizer_g = optim.Adam(self.growup.parameters(),
lr=1e-4,
betas=(0.5, 0.999))
self.optimizer_g2 = optim.Adam(self.growup.parameters(),
lr=1e-4,
betas=(0.5, 0.999))
self.optimizer_d = optim.Adam(self.discriminator.parameters(),
lr=1e-4,
betas=(0.5, 0.999))
self.optimizer2 = optim.SGD(
[{
'params': paras_wo_bn + [self.head.kernel],
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
if conf.finetune_model_path is not None:
self.optimizer = optim.SGD([{
'params': paras_wo_bn,
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}],
lr=conf.lr,
momentum=conf.momentum)
print('optimizers generated')
self.board_loss_every = len(self.loader) // 100
self.evaluate_every = len(self.loader) // 2
self.save_every = len(self.loader)
dataset_root = "/home/nas1_userD/yonggyu/Face_dataset/face_emore"
self.lfw = np.load(
os.path.join(dataset_root,
"lfw_align_112_list.npy")).astype(np.float32)
self.lfw_issame = np.load(
os.path.join(dataset_root, "lfw_align_112_label.npy"))
self.fgnetc = np.load(
os.path.join(dataset_root,
"FGNET_new_align_list.npy")).astype(np.float32)
self.fgnetc_issame = np.load(
os.path.join(dataset_root, "FGNET_new_align_label.npy"))
else:
# Will not use anymore
# self.model = nn.DataParallel(self.model)
self.threshold = conf.threshold
def __init__(self, conf, args, inference=False):
print(conf)
self.local_rank = args.local_rank
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).cuda()
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).cuda()
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.model, self.head], self.optimizer = amp.initialize([self.model, self.head], self.optimizer, opt_level='O1')
[self.model, self.head
], self.optimizer = amp.initialize([self.model, self.head],
self.optimizer,
opt_level='O3',
keep_batchnorm_fp32=True)
print(self.optimizer, args.local_rank)
self.head = DistributedDataParallel(self.head)
self.model = DistributedDataParallel(self.model)
#self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank])
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
print('optimizers generated')
self.board_loss_every = len(self.loader) // 100
self.evaluate_every = len(self.loader) // 10
self.save_every = len(self.loader) // 5
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
self.loader.dataset.root.parent)
else:
self.threshold = conf.threshold
def __init__(self,
conf,
inference=False,
train_transforms=None,
val_transforms=None,
train_loader=None):
print(conf)
if conf.use_mobilfacenet:
self.model = MobileFaceNet(conf.embedding_size).to(conf.device)
print('MobileFaceNet model generated')
else:
self.milestones = conf.milestones
if train_loader is None:
self.loader, self.class_num = get_train_loader(
conf, train_transforms)
else:
self.loader = train_loader
self.class_num = conf.num_classes
if conf.net_mode in ['ir', 'ir_se']:
self.model = Backbone(conf.net_depth, conf.drop_ratio,
conf.net_mode,
conf.use_gap).to(conf.device)
else:
import json
self.model = MetricNet(model_name=conf.net_mode,
pooling=conf.pooling,
use_fc=True,
fc_dim=conf.embedding_size,
dropout=conf.last_fc_dropout,
pretrained=conf.pretrained,
class_num=self.class_num).to(
conf.device)
print('{}_{} model generated'.format(conf.net_mode,
conf.net_depth))
if conf.use_mobilfacenet or conf.net_mode in ['ir', 'ir_se']:
self.head = Arcface(embedding_size=conf.embedding_size,
classnum=self.class_num).to(conf.device)
else:
if conf.loss_module == 'arcface':
self.head = ArcMarginProduct(self.model.final_in_features,
self.class_num,
s=conf.s,
m=conf.margin,
easy_margin=False,
ls_eps=conf.ls_eps).to(
conf.device)
elif conf.loss_module == 'cosface':
self.head = AddMarginProduct(self.model.final_in_features,
self.class_num,
s=conf.s,
m=conf.margin).to(conf.device)
elif conf.loss_module == 'adacos':
self.head = AdaCos(self.model.final_in_features,
self.class_num,
m=conf.margin,
theta_zero=conf.theta_zero).to(
conf.device)
else:
self.head = nn.Linear(self.model.final_in_features,
self.class_num).to(conf.device)
print('two model heads generated')
if conf.ft_model_path:
self.load_ft_model(conf.ft_model_path, not conf.no_strict)
elif conf.restore_suffix:
self.load_state(conf,
conf.restore_suffix,
from_save_folder=False,
model_only=False)
if not inference:
self.writer = SummaryWriter(conf.log_path)
self.step = 0
paras_only_bn, paras_wo_bn = separate_bn_paras(self.model)
if conf.use_mobilfacenet:
params = [{
'params': paras_wo_bn[:-1],
'weight_decay': 4e-5
}, {
'params': [paras_wo_bn[-1]] + [self.head.kernel],
'weight_decay': 4e-4
}, {
'params': paras_only_bn
}]
wd = 4e-5
else:
# if conf.net_mode in ['ir', 'ir_se']:
# params = [
# {'params': paras_wo_bn + [self.head.weight], 'weight_decay': 5e-4},
# {'params': paras_only_bn}
# ]
params = [{
'params': paras_wo_bn + [self.head.kernel],
'weight_decay': 5e-4
}, {
'params': paras_only_bn
}]
wd = 5e-4
# else:
# params = self.model.parameters()
# wd = conf.wd
# # params = [
# # {'params': paras_wo_bn + [self.head.weight], 'weight_decay': conf.wd}, # 5e-4},
# # {'params': paras_only_bn}
# # ]
if conf.optimizer == 'sgd':
self.optimizer = optim.SGD(
params, lr=conf.lr,
momentum=conf.momentum) # , weight_decay=wd)
elif conf.optimizer == 'adam':
self.optimizer = optim.Adam(
params, lr=conf.lr) # , weight_decay=wd)
print(self.optimizer)
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
print('optimizers generated')
self.board_loss_every = len(self.loader) // 100
self.evaluate_every = len(self.loader) // 10
self.save_every = len(self.loader) // 5
self.board_loss_every = 20
self.evaluate_every = len(self.loader)
self.save_every = len(self.loader)
if conf.data_mode == 'common':
import json
val_img_dir_map = json.loads(conf.val_img_dirs)
self.val_dataloaders = {}
for val_name in val_img_dir_map:
val_img_dir = val_img_dir_map[val_name]
val_dataloader, common_val_issame = get_common_val_data(
val_img_dir,
conf.max_positive_cnt,
conf.val_batch_size,
conf.val_pin_memory,
conf.num_workers,
val_transforms=val_transforms,
use_pos=not conf.not_use_pos,
use_neg=not conf.not_use_neg,
val_smapling_type=conf.val_smapling_type,
use_keras_model=conf.use_val_left_right_check)
self.val_dataloaders[val_name] = [
val_dataloader, common_val_issame
]
elif conf.data_mode == 'dacon_landmark':
pass
else:
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
self.loader.dataset.root.parent)
else:
self.threshold = conf.threshold
def main():
model_path = './softmax_loss_checkpoints/'
model = MobileFaceNet(conf.embedding_size)
model = nn.DataParallel(model)
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['net'])
from model import MobileFaceNet
import torch
from mtcnn import MTCNN
model = MobileFaceNet(512)
model.load_state_dict(torch.load('/home/ai/Desktop/project/pytorch-insightface/model_mobilefacenet.pth'))
print(model.state_dict())
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
except Exception as err:
print(err)
# create model
init_funcs = {
1: lambda x: torch.nn.init.normal_(x, mean=0., std=1.), # can be bias
2: lambda x: torch.nn.init.xavier_normal_(x, gain=1.), # can be weight
3: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.
), # can be conv1D filter
4: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.
), # can be conv2D filter
"default": lambda x: torch.nn.init.constant(x, 1.), # everything else
}
embedding_size = 512
model = MobileFaceNet(embedding_size)
init_all(model, init_funcs)
if mode == 'test':
load_model(model_name, model)
if cuda:
model = model.to(device)
if mode == 'train':
# define loss function
# loss_fn = nn.CrossEntropyLoss()
# if cuda:
# loss_fn = loss_fn.cuda()
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 __init__(self, conf, inference=False):
print(conf)
if conf.use_mobilfacenet:
self.model = MobileFaceNet(conf.embedding_size).to(conf.device)
print('MobileFaceNet model generated')
elif conf.use_shufflenetV2:
#self.model = ShuffleNetV2().to(conf.device)
self.model = model.shufflenet().to(conf.device)
print("ShufflenetV2 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)
# elif conf.use_shufflenetV2:
# 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*10, 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 * 10,
momentum=conf.momentum)
print(self.optimizer)
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
print('optimizers generated')
self.board_loss_every = len(self.loader) // 100
self.evaluate_every = len(self.loader) // 10
self.save_every = len(self.loader) // 5
self.agedb_30, self.cfp_fp, self.lfw, self.agedb_30_issame, self.cfp_fp_issame, self.lfw_issame = get_val_data(
'/root/NewDisk/daxing2/WW/data/faces_vgg_112x112/'
) #self.loader.dataset.root.parent
else:
self.threshold = conf.threshold
