global args, best_prec1

args = parser.parse_args()

conf_name = args.data

conf_name += '_ortho' if args.ortho else ''

conf_name += ('_pre_' + args.net_type) if args.net_type != 'default' else ''

args.checkpoint = os.path.join(args.checkpoint, conf_name)

if not os.path.isdir(args.checkpoint):

mkdir_p(args.checkpoint)

model = models.Net(extractor_type=args.net_type).cuda()

# define loss function (criterion) and optimizer

criterion = nn.CrossEntropyLoss().cuda()

extractor_params = list(map(id, model.extractor.parameters()))

classifier_params = filter(lambda p: id(p) not in extractor_params, model.parameters())

optimizer = torch.optim.SGD([

{'params': model.extractor.parameters()},

{'params': classifier_params, 'lr': args.lr * 10}

], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# optionally resume from a checkpoint

title = 'CUB'

if args.resume:

print('==> Resuming from checkpoint..')

assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'

checkpoint = torch.load(args.resume)

args.start_epoch = checkpoint['epoch']

best_prec1 = checkpoint['best_prec1']

model.load_state_dict(checkpoint['state_dict'])

optimizer.load_state_dict(checkpoint['optimizer'])

print("=> loaded checkpoint '{}' (epoch {})"

.format(args.resume, checkpoint['epoch']))

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.'])

cudnn.benchmark = True

print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))

# Data loading code

normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],

std=[0.5, 0.5, 0.5])

train_dataset = loader.ImageLoader(

args.data,

transforms.Compose([

transforms.RandomResizedCrop(224),

transforms.RandomHorizontalFlip(),

transforms.ToTensor(),

normalize,

]), train=True)

train_loader = torch.utils.data.DataLoader(

train_dataset, batch_size=args.batch_size, shuffle=True,

num_workers=args.workers, pin_memory=True)

val_loader = torch.utils.data.DataLoader(

loader.ImageLoader(args.data, transforms.Compose([

transforms.Resize(256),

transforms.CenterCrop(224),

transforms.ToTensor(),

normalize,

])),

batch_size=args.batch_size, shuffle=False,

num_workers=args.workers, pin_memory=True)

if args.evaluate:

validate(val_loader, model, criterion)

return

for epoch in range(args.start_epoch, args.epochs):

scheduler.step()

lr = optimizer.param_groups[1]['lr']

print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, lr))

# train for one epoch

train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch, args)

# evaluate on validation set

test_loss, test_acc = validate(val_loader, model, criterion)

# append logger file

logger.append([lr, train_loss, test_loss, train_acc, test_acc])

# remember best prec@1 and save checkpoint

is_best = test_acc > best_prec1

best_prec1 = max(test_acc, best_prec1)

save_checkpoint({

'epoch': epoch + 1,

'state_dict': model.state_dict(),

'best_prec1': best_prec1,

'optimizer' : optimizer.state_dict(),

}, is_best, checkpoint=args.checkpoint)

logger.close()

logger.plot()

savefig(os.path.join(args.checkpoint, 'log.eps'))

print('Best acc:')

batch_time = AverageMeter()

data_time = AverageMeter()

losses = AverageMeter()

top1 = AverageMeter()

top5 = AverageMeter()

# switch to train mode

model.train()

end = time.time()

bar = Bar('Training', max=len(train_loader))

for batch_idx, (input, target) in enumerate(train_loader):

# measure data loading time

data_time.update(time.time() - end)

input = input.cuda()

target = target.cuda(non_blocking=True)

# compute output

output = model(input)

loss = criterion(output, target)

if args.ortho:

weight_matrix = model.classifier.fc.weight

loss += args.ortho_lambda*ortho_reg.l2_reg_ortho(weight_matrix, 'cuda')

# measure accuracy and record loss

prec1, prec5 = accuracy(output, target, topk=(1, 5))

losses.update(loss.item(), input.size(0))

top1.update(prec1.item(), input.size(0))

top5.update(prec5.item(), input.size(0))

# compute gradient and do SGD step

optimizer.zero_grad()

loss.backward()

optimizer.step()

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

model.weight_norm()

# plot progress

bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(

batch=batch_idx + 1,

size=len(train_loader),

data=data_time.val,

bt=batch_time.val,

total=bar.elapsed_td,

eta=bar.eta_td,

loss=losses.avg,

top1=top1.avg,

top5=top5.avg,

)

bar.next()

bar.finish()

batch_time = AverageMeter()

data_time = AverageMeter()

losses = AverageMeter()

top1 = AverageMeter()

top5 = AverageMeter()

# switch to evaluate mode

model.eval()

bar = Bar('Testing ', max=len(val_loader))

with torch.no_grad():

end = time.time()

for batch_idx, (input, target) in enumerate(val_loader):

# measure data loading time

data_time.update(time.time() - end)

input = input.cuda()

target = target.cuda(non_blocking=True)

# compute output

output = model(input)

loss = criterion(output, target)

# measure accuracy and record loss

prec1, prec5 = accuracy(output, target, topk=(1, 5))

losses.update(loss.item(), input.size(0))

top1.update(prec1.item(), input.size(0))

top5.update(prec5.item(), input.size(0))

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

# plot progress

bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(

batch=batch_idx + 1,

size=len(val_loader),

data=data_time.avg,

bt=batch_time.avg,

total=bar.elapsed_td,

eta=bar.eta_td,

loss=losses.avg,

top1=top1.avg,

top5=top5.avg,

)

bar.next()

bar.finish()

filepath = os.path.join(checkpoint, filename)

torch.save(state, filepath)

if is_best: