def main():
parser = argparse.ArgumentParser(description='PyTorch Frame Attention Network Training')
parser.add_argument('--epochs', default=180, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--lr', '--learning-rate', default=4e-6, type=float,
metavar='LR', help='initial learning rate')
parser.add_argument('-e', '--evaluate', default=False, dest='evaluate', action='store_true',
help='evaluate model on validation set')
args = parser.parse_args()
best_acc = 0
logger = util.Logger('./log/','baseline_afew')
''' Load data '''
root_train = './data/face/train_afew'
list_train = './data/txt/afew_train.txt'
batchsize_train= 48
root_eval = './data/face/val_afew'
list_eval = './data/txt/afew_eval.txt'
batchsize_eval= 64
train_loader, val_loader = load.afew_faces_baseline(root_train, list_train, batchsize_train, root_eval, list_eval, batchsize_eval)
''' Load model '''
_structure = models.resnet18(num_classes=7)
_parameterDir = './pretrain_model/Resnet18_FER+_pytorch.pth.tar'
model = load.model_parameters(_structure, _parameterDir)
''' Loss & Optimizer '''
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=60, gamma=0.2)
cudnn.benchmark = True
''' Train & Eval '''
if args.evaluate == True:
logger.print('args.evaluate: {:}', args.evaluate)
validate(val_loader, model)
return
logger.print('baseline afew dataset, learning rate: {:}'.format(args.lr))
for epoch in range(args.epochs):
train(train_loader, model, optimizer, epoch, logger)
acc_epoch = val(val_loader, model, logger)
is_best = acc_epoch > best_acc
if is_best:
logger.print('better model!')
best_acc = max(acc_epoch, best_acc)
util.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'accuracy': acc_epoch,
}, at_type='baseline')
lr_scheduler.step()
logger.print("epoch: {:} learning rate:{:}".format(epoch+1, optimizer.param_groups[0]['lr']))
def main():
parser = argparse.ArgumentParser(description='PyTorch Frame Attention Network Training')
parser.add_argument('--at_type', '--attention', default=1, type=int, metavar='N',
help= '0 is self-attention; 1 is self + relation-attention')
parser.add_argument('--epochs', default=60, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('-f', '--fold', default=10, type=int, help='which fold used for ck+ test')
parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float,
metavar='LR', help='initial learning rate')
parser.add_argument('-e', '--evaluate', default=False, dest='evaluate', action='store_true',
help='evaluate model on validation set')
args = parser.parse_args()
best_acc = 0
at_type = ['self-attention', 'self_relation-attention'][args.at_type]
logger = util.Logger('./log/','fan_ckplus')
logger.print('The attention method is {:}, learning rate: {:}'.format(at_type, args.lr))
''' Load data '''
video_root = './data/face/ck_face'
video_list = './data/txt/CK+_10-fold_sample_IDascendorder_step10.txt'
batchsize_train= 48
batchsize_eval= 64
train_loader, val_loader = load.ckplus_faces_fan(video_root, video_list, args.fold, batchsize_train, batchsize_eval)
''' Load model '''
_structure = networks.resnet18_at(at_type=at_type)
_parameterDir = './pretrain_model/Resnet18_FER+_pytorch.pth.tar'
model = load.model_parameters(_structure, _parameterDir)
''' Loss & Optimizer '''
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.2)
cudnn.benchmark = True
''' Train & Eval '''
if args.evaluate == True:
logger.print('args.evaluate: {:}', args.evaluate)
validate(val_loader, model)
return
logger.print('frame attention network (fan) ck+ dataset, learning rate: {:}'.format(args.lr))
for epoch in range(args.epochs):
train(train_loader, model, optimizer, epoch)
acc_epoch = val(val_loader, model, at_type)
is_best = acc_epoch > best_acc
if is_best:
logger.print('better model!')
best_acc = max(acc_epoch, best_acc)
util.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'accuracy': acc_epoch,
}, at_type=at_type)
lr_scheduler.step()
logger.print("epoch: {:} learning rate:{:}".format(epoch+1, optimizer.param_groups[0]['lr']))