def build_model(args):
# model = ResNet32(args.dataset == 'cifar10' and 10 or 100)
model = WideResNet(args.layers,
args.dataset == 'cifar10' and 10 or 100,
args.widen_factor,
dropRate=args.droprate)
# weights_init(model)
# print('Number of model parameters: {}'.format(
# sum([p.data.nelement() for p in model.params()])))
if torch.cuda.is_available():
model.cuda()
torch.backends.cudnn.benchmark = True
return model
cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10).cuda()
elif args.model == 'wideresnet16_8':
cnn = WideResNet(depth=16, num_classes=num_classes, widen_factor=8).cuda()
elif args.model == 'densenet':
cnn = DenseNet3(depth=100,
num_classes=num_classes,
growth_rate=12,
reduction=0.5).cuda()
elif args.model == 'vgg13':
cnn = VGG(vgg_name='VGG13', num_classes=num_classes).cuda()
model_dict = cnn.state_dict()
pretrained_dict = torch.load('checkpoints/' + filename + '.pt')
cnn.load_state_dict(pretrained_dict)
cnn = cnn.cuda()
cnn.eval()
#------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# Evaluate関数
#------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
def evaluate(data_loader, mode):
out = []
xent = nn.CrossEntropyLoss()
for data in data_loader:
if type(data) == list:
images, labels = data
else:
def main():
global args, best_prec1
args = parser.parse_args()
# torch.cuda.set_device(args.gpu)
if args.tensorboard:
print("Using tensorboard")
configure("exp/%s" % (args.name))
# Data loading code
if args.augment:
print(
"Doing image augmentation with\n"
"Zoom: prob: {zoom_prob} range: {zoom_range}\n"
"Stretch: prob: {stretch_prob} range: {stretch_range}\n"
"Rotation: prob: {rotation_prob} range: {rotation_degree}".format(
zoom_prob=args.zoom_prob,
zoom_range=args.zoom_range,
stretch_prob=args.stretch_prob,
stretch_range=args.stretch_range,
rotation_prob=args.rotation_prob,
rotation_degree=args.rotation_degree))
transform_train = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: F.pad(
Variable(x.unsqueeze(0), requires_grad=False, volatile=True),
(4, 4, 4, 4),
mode='replicate').data.squeeze()),
transforms.ToPILImage(),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(prob=args.rotation_prob,
degree=args.rotation_degree),
transforms.RandomZoom(prob=args.zoom_prob,
zoom_range=args.zoom_range),
transforms.RandomStretch(prob=args.stretch_prob,
stretch_range=args.stretch_range),
transforms.ToTensor(),
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
kwargs = {'num_workers': 1, 'pin_memory': True}
assert (args.dataset == 'cifar10' or args.dataset == 'cifar100')
train_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data',
train=True,
download=True,
transform=transform_train),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data',
train=False,
transform=transform_test),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
# create model
model = WideResNet(args.layers,
args.dataset == 'cifar10' and 10 or 100,
args.widen_factor,
dropRate=args.droprate)
# get the number of model parameters
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
# for training on multiple GPUs.
# Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
# optionally 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_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch + 1)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
print 'Best accuracy: ', best_prec1