def get_train_transforms(normalize): train_transforms = [] train_transforms.append(transforms.Scale(160)) train_transforms.append(transforms.RandomHorizontalFlip()) train_transforms.append(transforms.RandomColor(0.15)) train_transforms.append(transforms.RandomRotate(15)) train_transforms.append(transforms.RandomSizedCrop(128)) train_transforms.append(transforms.ToTensor()) train_transforms.append(normalize) train_transforms = transforms.Compose(train_transforms) return train_transforms
def main(args):
# parse args
best_acc1 = 0.0
if args.gpu >= 0:
print("Use GPU: {}".format(args.gpu))
else:
print('You are using CPU for computing!',
'Yet we assume you are using a GPU.',
'You will NOT be able to switch between CPU and GPU training!')
# set up the model + loss
if args.use_custom_conv:
print("Using custom convolutions in the network")
model = default_model(conv_op=CustomConv2d, num_classes=100)
else:
model = default_model(num_classes=100)
model_arch = "simplenet"
# model_arch = "simplenet_batchnorm2d"
# model_arch = "resnet18"
# model_arch = "resnet34"
criterion = nn.CrossEntropyLoss()
# put everthing to gpu
if args.gpu >= 0:
model = model.cuda(args.gpu)
criterion = criterion.cuda(args.gpu)
# setup the optimizer
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# resume from a checkpoint?
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
if args.gpu < 0:
model = model.cpu()
else:
model = model.cuda(args.gpu)
# only load the optimizer if necessary
if (not args.evaluate) and (not args.attack):
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {}, acc1 {})"
.format(args.resume, checkpoint['epoch'], best_acc1))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# set up transforms for data augmentation
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transforms = []
train_transforms.append(transforms.Scale(160))
train_transforms.append(transforms.RandomHorizontalFlip())
train_transforms.append(transforms.RandomColor(0.15))
train_transforms.append(transforms.RandomRotate(15))
train_transforms.append(transforms.RandomSizedCrop(128))
train_transforms.append(transforms.ToTensor())
train_transforms.append(normalize)
train_transforms = transforms.Compose(train_transforms)
# val transofrms
val_transforms=[]
val_transforms.append(transforms.Scale(160, interpolations=None))
val_transforms.append(transforms.ToTensor())
val_transforms.append(normalize)
val_transforms = transforms.Compose(val_transforms)
if (not args.evaluate) and (not args.attack):
print("Training time data augmentations:")
print(train_transforms)
# setup dataset and dataloader
train_dataset = MiniPlacesLoader(args.data_folder,
split='train', transforms=train_transforms)
val_dataset = MiniPlacesLoader(args.data_folder,
split='val', transforms=val_transforms)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, sampler=None, drop_last=True)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=100, shuffle=False,
num_workers=args.workers, pin_memory=True, sampler=None, drop_last=False)
# testing only
if (args.evaluate==args.attack) and args.evaluate:
print("Cann't set evaluate and attack to True at the same time!")
return
# set up visualizer
if args.vis:
visualizer = default_attention(criterion)
else:
visualizer = None
# evaluation
if args.resume and args.evaluate:
print("Testing the model ...")
cudnn.deterministic = True
validate(val_loader, model, -1, args, visualizer=visualizer)
return
# attack
if args.resume and args.attack:
print("Generating adversarial samples for the model ..")
cudnn.deterministic = True
validate(val_loader, model, -1, args,
attacker=default_attack(criterion),
visualizer=visualizer)
return
# enable cudnn benchmark
cudnn.enabled = True
cudnn.benchmark = True
# warmup the training
if (args.start_epoch == 0) and (args.warmup_epochs > 0):
print("Warmup the training ...")
for epoch in range(0, args.warmup_epochs):
train(train_loader, model, criterion, optimizer, epoch, "warmup", args)
# start the training
print("Training the model ...")
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, "train", args)
# evaluate on validation set
acc1 = validate(val_loader, model, epoch, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
save_checkpoint({
'epoch': epoch + 1,
'model_arch': model_arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer' : optimizer.state_dict(),
}, is_best)