def train(args):
# gpu init
multi_gpus = False
if len(args.gpus.split(',')) > 1:
multi_gpus = True
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# log init
save_dir = os.path.join(
args.save_dir, args.model_pre + args.backbone.upper() + '_' +
datetime.now().strftime('%Y%m%d_%H%M%S'))
if os.path.exists(save_dir):
raise NameError('model dir exists!')
os.makedirs(save_dir)
logging = init_log(save_dir)
_print = logging.info
# dataset loader
transform = transforms.Compose([
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5,
0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])
# validation dataset
trainset = CASIAWebFace(args.train_root,
args.train_file_list,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
drop_last=False)
# test dataset
lfwdataset = LFW(args.lfw_test_root,
args.lfw_file_list,
transform=transform)
lfwloader = torch.utils.data.DataLoader(lfwdataset,
batch_size=128,
shuffle=False,
num_workers=4,
drop_last=False)
agedbdataset = AgeDB30(args.agedb_test_root,
args.agedb_file_list,
transform=transform)
agedbloader = torch.utils.data.DataLoader(agedbdataset,
batch_size=128,
shuffle=False,
num_workers=4,
drop_last=False)
cfpfpdataset = CFP_FP(args.cfpfp_test_root,
args.cfpfp_file_list,
transform=transform)
cfpfploader = torch.utils.data.DataLoader(cfpfpdataset,
batch_size=128,
shuffle=False,
num_workers=4,
drop_last=False)
# define backbone and margin layer
if args.backbone == 'MobileFace':
net = MobileFaceNet()
elif args.backbone == 'Res50':
net = ResNet50()
elif args.backbone == 'Res101':
net = ResNet101()
elif args.backbone == 'Res50_IR':
net = SEResNet_IR(50, feature_dim=args.feature_dim, mode='ir')
elif args.backbone == 'SERes50_IR':
net = SEResNet_IR(50, feature_dim=args.feature_dim, mode='se_ir')
elif args.backbone == 'SphereNet':
net = SphereNet(num_layers=64, feature_dim=args.feature_dim)
else:
print(args.backbone, ' is not available!')
if args.margin_type == 'ArcFace':
margin = ArcMarginProduct(args.feature_dim,
trainset.class_nums,
s=args.scale_size)
elif args.margin_type == 'CosFace':
pass
elif args.margin_type == 'SphereFace':
pass
elif args.margin_type == 'InnerProduct':
margin = InnerProduct(args.feature_dim, trainset.class_nums)
else:
print(args.margin_type, 'is not available!')
if args.resume:
print('resume the model parameters from: ', args.net_path,
args.margin_path)
net.load_state_dict(torch.load(args.net_path)['net_state_dict'])
margin.load_state_dict(torch.load(args.margin_path)['net_state_dict'])
# define optimizers for different layers
criterion_classi = torch.nn.CrossEntropyLoss().to(device)
optimizer_classi = optim.SGD([{
'params': net.parameters(),
'weight_decay': 5e-4
}, {
'params': margin.parameters(),
'weight_decay': 5e-4
}],
lr=0.1,
momentum=0.9,
nesterov=True)
#criterion_center = CenterLoss(trainset.class_nums, args.feature_dim).to(device)
#optimizer_center = optim.SGD(criterion_center.parameters(), lr=0.5)
scheduler_classi = lr_scheduler.MultiStepLR(optimizer_classi,
milestones=[25, 50, 65],
gamma=0.1)
if multi_gpus:
net = DataParallel(net).to(device)
margin = DataParallel(margin).to(device)
else:
net = net.to(device)
margin = margin.to(device)
best_lfw_acc = 0.0
best_lfw_iters = 0
best_agedb30_acc = 0.0
best_agedb30_iters = 0
best_cfp_fp_acc = 0.0
best_cfp_fp_iters = 0
total_iters = 0
#vis = Visualizer(env='softmax_center_xavier')
for epoch in range(1, args.total_epoch + 1):
scheduler_classi.step()
# train model
_print('Train Epoch: {}/{} ...'.format(epoch, args.total_epoch))
net.train()
since = time.time()
for data in trainloader:
img, label = data[0].to(device), data[1].to(device)
feature = net(img)
output = margin(feature)
loss_classi = criterion_classi(output, label)
#loss_center = criterion_center(feature, label)
total_loss = loss_classi #+ loss_center * args.weight_center
optimizer_classi.zero_grad()
#optimizer_center.zero_grad()
total_loss.backward()
optimizer_classi.step()
#optimizer_center.step()
total_iters += 1
# print train information
if total_iters % 100 == 0:
# current training accuracy
_, predict = torch.max(output.data, 1)
total = label.size(0)
correct = (np.array(predict) == np.array(label.data)).sum()
time_cur = (time.time() - since) / 100
since = time.time()
#vis.plot_curves({'softmax loss': loss_classi.item(), 'center loss': loss_center.item()}, iters=total_iters, title='train loss', xlabel='iters', ylabel='train loss')
#vis.plot_curves({'train accuracy': correct / total}, iters=total_iters, title='train accuracy', xlabel='iters', ylabel='train accuracy')
print(
"Iters: {:0>6d}/[{:0>2d}], loss_classi: {:.4f}, loss_center: {:.4f}, train_accuracy: {:.4f}, time: {:.2f} s/iter, learning rate: {}"
.format(total_iters, epoch, loss_classi.item(),
loss_center.item(), correct / total, time_cur,
scheduler_classi.get_lr()[0]))
# save model
if total_iters % args.save_freq == 0:
msg = 'Saving checkpoint: {}'.format(total_iters)
_print(msg)
if multi_gpus:
net_state_dict = net.module.state_dict()
margin_state_dict = margin.module.state_dict()
else:
net_state_dict = net.state_dict()
margin_state_dict = margin.state_dict()
if not os.path.exists(save_dir):
os.mkdir(save_dir)
torch.save(
{
'iters': total_iters,
'net_state_dict': net_state_dict
},
os.path.join(save_dir, 'Iter_%06d_net.ckpt' % total_iters))
torch.save(
{
'iters': total_iters,
'net_state_dict': margin_state_dict
},
os.path.join(save_dir,
'Iter_%06d_margin.ckpt' % total_iters))
#torch.save({
# 'iters': total_iters,
# 'net_state_dict': criterion_center.state_dict()},
# os.path.join(save_dir, 'Iter_%06d_center.ckpt' % total_iters))
# test accuracy
if total_iters % args.test_freq == 0:
# test model on lfw
net.eval()
getFeatureFromTorch('./result/cur_lfw_result.mat', net, device,
lfwdataset, lfwloader)
lfw_accs = evaluation_10_fold('./result/cur_lfw_result.mat')
_print('LFW Ave Accuracy: {:.4f}'.format(
np.mean(lfw_accs) * 100))
if best_lfw_acc < np.mean(lfw_accs) * 100:
best_lfw_acc = np.mean(lfw_accs) * 100
best_lfw_iters = total_iters
# test model on AgeDB30
getFeatureFromTorch('./result/cur_agedb30_result.mat', net,
device, agedbdataset, agedbloader)
age_accs = evaluation_10_fold(
'./result/cur_agedb30_result.mat')
_print('AgeDB-30 Ave Accuracy: {:.4f}'.format(
np.mean(age_accs) * 100))
if best_agedb30_acc < np.mean(age_accs) * 100:
best_agedb30_acc = np.mean(age_accs) * 100
best_agedb30_iters = total_iters
# test model on CFP-FP
getFeatureFromTorch('./result/cur_cfpfp_result.mat', net,
device, cfpfpdataset, cfpfploader)
cfp_accs = evaluation_10_fold('./result/cur_cfpfp_result.mat')
_print('CFP-FP Ave Accuracy: {:.4f}'.format(
np.mean(cfp_accs) * 100))
if best_cfp_fp_acc < np.mean(cfp_accs) * 100:
best_cfp_fp_acc = np.mean(cfp_accs) * 100
best_cfp_fp_iters = total_iters
_print(
'Current Best Accuracy: LFW: {:.4f} in iters: {}, AgeDB-30: {:.4f} in iters: {} and CFP-FP: {:.4f} in iters: {}'
.format(best_lfw_acc, best_lfw_iters, best_agedb30_acc,
best_agedb30_iters, best_cfp_fp_acc,
best_cfp_fp_iters))
#vis.plot_curves({'lfw': np.mean(lfw_accs), 'agedb-30': np.mean(age_accs), 'cfp-fp': np.mean(cfp_accs)}, iters=total_iters,
# title='test accuracy', xlabel='iters', ylabel='test accuracy')
net.train()
_print(
'Finally Best Accuracy: LFW: {:.4f} in iters: {}, AgeDB-30: {:.4f} in iters: {} and CFP-FP: {:.4f} in iters: {}'
.format(best_lfw_acc, best_lfw_iters, best_agedb30_acc,
best_agedb30_iters, best_cfp_fp_acc, best_cfp_fp_iters))
print('finishing training')
def train(args):
# gpu init
multi_gpus = False
if len(args.gpus.split(',')) > 1:
multi_gpus = True
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# log init
save_dir = os.path.join(
args.save_dir, args.model_pre + args.backbone.upper() + '_' +
datetime.now().strftime('%Y%m%d_%H%M%S'))
if os.path.exists(save_dir):
raise NameError('model dir exists!')
os.makedirs(save_dir)
logging = init_log(save_dir)
_print = logging.info
# dataset loader
if not args.use_gray:
transform = transforms.Compose([
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5,
0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])
else:
transform = transforms.Compose([
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
])
# validation dataset
trainset = CASIAWebFace(args.train_root,
args.train_file_list,
transform=transform,
use_gray=args.use_gray)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
drop_last=False)
# test dataset
lfwdataset = LFW(args.lfw_test_root,
args.lfw_file_list,
transform=transform,
use_gray=args.use_gray)
lfwloader = torch.utils.data.DataLoader(lfwdataset,
batch_size=128,
shuffle=False,
num_workers=4,
drop_last=False)
# define backbone and margin layer
in_channels = 1 if args.use_gray else 3
if args.backbone == 'MobileFace':
net = MobileFaceNet()
elif args.backbone == 'Res50_IR':
net = CBAMResNet(50, feature_dim=args.feature_dim, mode='ir')
elif args.backbone == 'SERes50_IR':
net = CBAMResNet(50, feature_dim=args.feature_dim, mode='ir_se')
elif args.backbone == 'Res100_IR':
net = CBAMResNet(100, feature_dim=args.feature_dim, mode='ir')
elif args.backbone == 'SERes100_IR':
net = CBAMResNet(100, feature_dim=args.feature_dim, mode='ir_se')
elif args.backbone == 'Attention_56':
net = ResidualAttentionNet_56(feature_dim=args.feature_dim)
elif args.backbone == 'Attention_92':
net = ResidualAttentionNet_92(feature_dim=args.feature_dim)
elif args.backbone == 'SmallVGG':
net = SmallVGG(in_channels, args.feature_dim, alpha=0.5)
else:
print(args.backbone, ' is not available!')
exit(-1)
calc_flops(net, in_channels, 112, 112)
if args.margin_type == 'ArcFace':
margin = ArcMarginProduct(args.feature_dim,
trainset.class_nums,
s=args.scale_size)
elif args.margin_type == 'MultiMargin':
margin = MultiMarginProduct(args.feature_dim,
trainset.class_nums,
s=args.scale_size)
elif args.margin_type == 'CosFace':
margin = CosineMarginProduct(args.feature_dim,
trainset.class_nums,
s=args.scale_size)
elif args.margin_type == 'Softmax':
margin = InnerProduct(args.feature_dim, trainset.class_nums)
elif args.margin_type == 'SphereFace':
pass
else:
print(args.margin_type, 'is not available!')
if args.resume:
print('resume the model parameters from: ', args.net_path,
args.margin_path)
net.load_state_dict(torch.load(args.net_path)['net_state_dict'])
margin.load_state_dict(torch.load(args.margin_path)['net_state_dict'])
# define optimizers for different layer
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer_ft = optim.SGD([{
'params': net.parameters(),
'weight_decay': 5e-4
}, {
'params': margin.parameters(),
'weight_decay': 5e-4
}],
lr=0.1,
momentum=0.9,
nesterov=True)
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft,
milestones=[6, 11, 16],
gamma=0.1)
if multi_gpus:
net = DataParallel(net).to(device)
margin = DataParallel(margin).to(device)
else:
net = net.to(device)
margin = margin.to(device)
best_lfw_acc = 0.0
best_lfw_iters = 0
total_iters = 0
for epoch in range(1, args.total_epoch + 1):
if epoch > 1:
exp_lr_scheduler.step()
# train model
_print('Train Epoch: {}/{} ...'.format(epoch, args.total_epoch))
net.train()
since = time.time()
for data in trainloader:
img, label = data[0].to(device), data[1].to(device)
optimizer_ft.zero_grad()
raw_logits = net(img)
output = margin(raw_logits, label)
total_loss = criterion(output, label)
total_loss.backward()
optimizer_ft.step()
total_iters += 1
# print train information
if total_iters % 100 == 0:
# current training accuracy
_, predict = torch.max(output.data, 1)
total = label.size(0)
correct = (np.array(predict.cpu()) == np.array(
label.data.cpu())).sum()
time_cur = (time.time() - since) / 100
since = time.time()
_print(
"Iters: {:0>6d}/[{:0>2d}], loss: {:.4f}, train_accuracy: {:.4f}, time: {:.2f} s/iter, learning rate: {}"
.format(total_iters, epoch, total_loss.item(),
correct / total, time_cur,
exp_lr_scheduler.get_last_lr()[0]))
# save model
if total_iters % args.save_freq == 0:
msg = 'Saving checkpoint: {}'.format(total_iters)
_print(msg)
if multi_gpus:
net_state_dict = net.module.state_dict()
margin_state_dict = margin.module.state_dict()
else:
net_state_dict = net.state_dict()
margin_state_dict = margin.state_dict()
if not os.path.exists(save_dir):
os.mkdir(save_dir)
torch.save(
{
'iters': total_iters,
'net_state_dict': net_state_dict
},
os.path.join(save_dir, 'Iter_%06d_net.ckpt' % total_iters))
torch.save(
{
'iters': total_iters,
'net_state_dict': margin_state_dict
},
os.path.join(save_dir,
'Iter_%06d_margin.ckpt' % total_iters))
# test accuracy
if total_iters % args.test_freq == 0:
# test model on lfw
net.eval()
getFeatureFromTorch('./result/cur_lfw_result.mat', net, device,
lfwdataset, lfwloader)
lfw_accs = evaluation_10_fold('./result/cur_lfw_result.mat')
_print('LFW Ave Accuracy: {:.4f}'.format(
np.mean(lfw_accs) * 100))
if best_lfw_acc <= np.mean(lfw_accs) * 100:
best_lfw_acc = np.mean(lfw_accs) * 100
best_lfw_iters = total_iters
net.train()
_print('Finally Best Accuracy: LFW: {:.4f} in iters: {}'.format(
best_lfw_acc, best_lfw_iters))
print('finishing training')