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)