def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
print('==> Preparing dataset %s' % args.dataset)
if args.dataset == 'cifar100':
training_loader = get_training_dataloader(settings.CIFAR100_TRAIN_MEAN,
settings.CIFAR100_TRAIN_STD,
num_workers=4,
batch_size=args.train_batch,
shuffle=True)
test_loader = get_test_dataloader(settings.CIFAR100_TRAIN_MEAN,
settings.CIFAR100_TRAIN_STD,
num_workers=4,
batch_size=args.test_batch,
shuffle=False)
num_classes = 100
else:
training_loader = get_training_dataloader_10(
settings.CIFAR10_TRAIN_MEAN,
settings.CIFAR10_TRAIN_STD,
num_workers=4,
batch_size=args.train_batch,
shuffle=True)
test_loader = get_test_dataloader_10(settings.CIFAR10_TRAIN_MEAN,
settings.CIFAR10_TRAIN_STD,
num_workers=4,
batch_size=args.test_batch,
shuffle=False)
num_classes = 10
#data preprocessing:
print("==> creating model '{}'".format(args.arch))
model = get_network(args, num_classes=num_classes)
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion1 = am_softmax.AMSoftmax()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
title = 'cifar-10-' + args.arch
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
train_scheduler = optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.schedule, gamma=0.2) #learning rate decay
iter_per_epoch = len(training_loader)
warmup_scheduler = WarmUpLR(optimizer, iter_per_epoch * args.warm)
for epoch in range(start_epoch, args.epochs):
if epoch > args.warm:
train_scheduler.step(epoch)
train_loss, train_acc = train(training_loader, model, warmup_scheduler,
criterion, criterion1, optimizer, epoch,
use_cuda)
test_loss, test_acc = eval_training(test_loader, model, criterion,
epoch, use_cuda)
logger.append([
optimizer.param_groups[0]['lr'], train_loss, test_loss, train_acc,
test_acc
])
# save model
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
checkpoint=args.checkpoint)
logger.close()
# logger.plot()
# savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
