output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
label = label.data.cpu().numpy()
# print(output)
# print(label)
acc = np.mean((output == label).astype(int))
spend_time = (time.time() - start)
train_acc.extend([acc] * len(label))
train_loss.extend([loss.item()] * len(label))
train_spend.extend([spend_time])
time_str = time.asctime(time.localtime(time.time()))
start = time.time()
print('{} train epoch {} {} seconds/epoch loss {} acc {}'.format(
time_str, i, np.sum(train_spend), np.mean(train_loss), np.mean(train_acc)))
if opt.display:
visualizer.display_current_results(i, np.mean(train_loss), name='train_loss')
visualizer.display_current_results(i, np.mean(train_acc), name='train_acc')
scheduler.step()
if i % opt.save_interval == 0 or i == opt.max_epoch:
save_model(model, opt.checkpoints_path, opt.backbone, i)
model.eval()
eval_acc = []
eval_loss = []
eval_speed = []
for ii, data in enumerate(val_loader):
data_input, label = data
data_input = data_input.to(device)
label = label.to(device).long()
def main(args):
opt = Config()
opt.num_classes = len(get_train_labels(opt.train_root, opt.criteria_list))
if opt.display:
visualizer = Visualizer()
if torch.cuda.is_available(): # GPUが利用可能か確認
device = 'cuda'
else:
device = 'cpu'
print('device: {}'.format(device))
train_dataset = Dataset(opt.train_root,
opt.train_list,
phase='train',
input_shape=opt.input_shape)
trainloader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
val_dataset = Dataset(opt.train_root,
opt.val_list,
phase='val',
input_shape=opt.input_shape)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.test_batch_size,
shuffle=True,
num_workers=opt.num_workers)
print('{} train iters per epoch:'.format(len(trainloader)))
if opt.loss == 'focal_loss':
criterion = FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet18':
model = resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet34':
model = resnet34()
elif opt.backbone == 'resnet50':
model = resnet50()
else:
raise TypeError('not match model type')
model.to(device)
model = DataParallel(model)
if args.train_second:
opt.metric = 'liner'
if device == 'cuda':
model.load_state_dict(torch.load(opt.test_model_path))
else:
model.load_state_dict(
torch.load(opt.test_model_path, map_location={'cuda:0':
'cpu'}))
for param in model.parameters():
param.requires_grad = False
model.eval()
print(model)
if opt.metric == 'add_margin':
metric_fc = AddMarginProduct(512, opt.num_classes, s=30, m=0.35)
elif opt.metric == 'arc_margin':
metric_fc = ArcMarginProduct(512,
opt.num_classes,
s=30,
m=0.5,
easy_margin=opt.easy_margin)
elif opt.metric == 'sphere':
metric_fc = SphereProduct(512, opt.num_classes, m=4)
else:
metric_fc = nn.Linear(512, opt.num_classes)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
if opt.optimizer == 'sgd':
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
else:
optimizer = torch.optim.Adam([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
scheduler = StepLR(optimizer, step_size=opt.lr_step, gamma=0.1)
start = time.time()
for i in range(opt.max_epoch):
model.train()
train_acc = []
train_loss = []
train_spend = []
for ii, train_batch in enumerate(trainloader):
data_input, label = train_batch
data_input = data_input.to(device)
label = label.to(device).long()
feature = model(data_input)
if args.train_second:
output = metric_fc(feature)
else:
output = metric_fc(feature, label)
loss = criterion(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
iters = i * len(trainloader) + ii
output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
label = label.data.cpu().numpy()
acc = np.mean((output == label).astype(int))
spend_time = (time.time() - start)
train_acc.extend([acc] * len(label))
train_loss.extend([loss.item()] * len(label))
train_spend.extend([spend_time])
time_str = time.asctime(time.localtime(time.time()))
start = time.time()
print('{} train epoch {} {} seconds/epoch loss {} acc {}'.format(
time_str, i, np.sum(train_spend), np.mean(train_loss),
np.mean(train_acc)))
if opt.display:
visualizer.display_current_results(i,
np.mean(train_loss),
name='train_loss')
visualizer.display_current_results(i,
np.mean(train_acc),
name='train_acc')
scheduler.step()
if args.train_second:
save_model(metric_fc, opt.checkpoints_path, 'fc', i)
else:
if i % opt.save_interval == 0 or i == opt.max_epoch:
save_model(model, opt.checkpoints_path, opt.backbone, i)
save_model(metric_fc, opt.checkpoints_path, 'metric_fc', i)
model.eval()
eval_acc = []
eval_loss = []
eval_speed = []
for ii, val_batch in enumerate(val_loader):
data_input, label = val_batch
data_input = data_input.to(device)
label = label.to(device).long()
feature = model(data_input)
if args.train_second:
output = metric_fc(feature)
else:
output = metric_fc(feature, label)
loss = criterion(output, label)
output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
label = label.data.cpu().numpy()
acc = np.mean((output == label).astype(int))
speed = opt.print_freq / (time.time() - start)
eval_acc.extend([acc] * len(label))
eval_loss.extend([loss.item()] * len(label))
eval_speed.extend([speed] * len(label))
time_str = time.asctime(time.localtime(time.time()))
print('{} val epoch {} loss {} acc {}'.format(time_str, i,
np.mean(eval_speed),
np.mean(eval_loss),
np.mean(eval_acc)))
if opt.display:
visualizer.display_current_results(i,
np.mean(eval_loss),
name='val_loss')
visualizer.display_current_results(i,
np.mean(eval_acc),
name='val_acc')
loss_all = loss_fn(output, label)
loss = loss_all.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
iters = i * len(trainloader) + ii
if iters % opt.print_freq == 0:
output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
label = label.data.cpu().numpy()
acc = np.mean((output == label).astype(int))
speed = opt.print_freq / (time.time() - start)
time_str = time.asctime(time.localtime(time.time()))
print('{} train epoch {} iter {} {} iters/s loss {} acc {}'.format(time_str, i, ii, speed, loss.item(), acc))
if opt.display:
visualizer.display_current_results(iters, loss.item(), name='train_loss')
visualizer.display_current_results(iters, acc, name='train_acc')
start = time.time()
if i % opt.save_interval == 0 or i == opt.max_epoch:
save_model(model, opt.checkpoints_path, opt.backbone, i)
model.eval()
loss = lfw_test(model, img_paths, identity_list, opt.lfw_test_list, opt.test_batch_size)
if opt.display:
visualizer.display_current_results(iters, loss, name='val_loss')
optimizer.step()
iters = i * len(trainloader) + ii
if iters % opt.print_freq == 0:
output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
label = label.data.cpu().numpy()
acc = np.mean((output == label).astype(int))
speed = opt.print_freq / (time.time() - start)
time_str = time.asctime(time.localtime(time.time()))
print('{} train epoch {} iter {} {} iters/s loss {} acc {}'.
format(time_str, i, ii, speed, loss.item(), acc))
if opt.display:
visualizer.display_current_results(iters,
loss.item(),
name='train_loss')
visualizer.display_current_results(iters,
acc,
name='train_acc')
start = time.time()
if i % opt.save_interval == 0 or i == opt.max_epoch:
save_model(model, opt.checkpoints_path, opt.backbone, i)
model.eval()
acc, result = lfw_test(model, img_paths, identity_list,
opt.lfw_test_list, opt.test_batch_size)
if opt.display:
