def inference_program(model_name, num_class):
# The image is 32 * 32 with RGB representation.
data_shape = [3, 32, 32]
images = fluid.layers.data(name='pixel', shape=data_shape, dtype='float32')
if model_name == 'ResNet20':
predict = resnet_cifar(images, 20, num_class)
elif model_name == 'ResNet32':
predict = resnet_cifar(images, 32, num_class)
elif model_name == 'ResNet110':
predict = resnet_cifar(images, 110, num_class)
else:
predict = NASCifarNet(images, 36, 6, 3, num_class,
Networks[model_name], True)
return predict
model.load_state_dict(best_model)
return model,cost_time,best_acc,best_train_acc
if __name__ == '__main__':
print ('DataSets: '+args.dataset)
print ('ResNet Depth: '+str(args.depth))
loader = DataLoader(args.dataset,batch_size=args.batch_size)
dataloaders,dataset_sizes = loader.load_data()
num_classes = 10
if args.dataset == 'cifar-10':
num_classes = 10
if args.dataset == 'cifar-100':
num_classes = 100
model = resnet_cifar(depth=args.depth, num_classes=num_classes)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1,
momentum=0.9, nesterov=True, weight_decay=1e-4)
# define loss and optimizer
criterion = nn.CrossEntropyLoss()
scheduler = MultiStepLR(optimizer, milestones=[args.epoch*0.4, args.epoch*0.6, args.epoch*0.8], gamma=0.1)
use_gpu = torch.cuda.is_available()
if use_gpu:
try:
args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]
except ValueError:
raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only')
model = torch.nn.DataParallel(model, device_ids=args.gpu_ids)
# patch_replication_callback(model)