def load_network(loc, masked=False):
net_checkpoint = torch.load(loc)
start_epoch = net_checkpoint['epoch']
SavedConv, SavedBlock = what_conv_block(net_checkpoint['conv'],
net_checkpoint['blocktype'],
net_checkpoint['module'])
net = WideResNet(args.wrn_depth,
args.wrn_width,
SavedConv,
SavedBlock,
num_classes=num_classes,
dropRate=0,
masked=masked).cuda()
if masked:
new_sd = net.state_dict()
old_sd = net_checkpoint['net']
new_names = [v for v in new_sd]
old_names = [v for v in old_sd]
for i, j in enumerate(new_names):
new_sd[j] = old_sd[old_names[i]]
net.load_state_dict(new_sd)
else:
net.load_state_dict(net_checkpoint['net'])
return net, start_epoch
def load_network(loc):
net_checkpoint = torch.load(loc)
start_epoch = net_checkpoint['epoch']
SavedConv, SavedBlock = what_conv_block(net_checkpoint['conv'],
net_checkpoint['blocktype'], net_checkpoint['module'])
net = WideResNet(args.wrn_depth, args.wrn_width, SavedConv, SavedBlock, num_classes=num_classes, dropRate=0).cuda()
net.load_state_dict(net_checkpoint['net'])
return net, start_epoch
def load_baseline_model(self):
"""
Load a simple baseline model AND dataset
Note that this sets the model to training mode
"""
if self.args.dataset == DATASET_CIFAR_10:
imsize, in_channel, num_classes = 32, 3, 10
elif self.args.dataset == DATASET_CIFAR_100:
imsize, in_channel, num_classes = 32, 3, 100
elif self.args.dataset == DATASET_MNIST:
imsize, in_channel, num_classes = 28, 1, 10
elif self.args.dataset == DATASET_BOSTON:
imsize, in_channel, num_classes = 13, 1, 1
# init_l2 = -7 # TODO: Important to make sure this is small enough to be unregularized when starting?
if self.args.model == MODEL_RESNET18:
cnn = ResNet18(num_classes=num_classes)
elif self.args.model == MODEL_WIDERESNET:
cnn = WideResNet(depth=28,
num_classes=num_classes,
widen_factor=10,
dropRate=0.3)
elif self.args.model[:3] == MODEL_MLP:
cnn = Net(self.args.num_layers,
0.0,
imsize,
in_channel,
INIT_L2,
num_classes=num_classes,
do_classification=self.args.do_classification)
elif self.args.model == MODEL_CNN_MLP:
cnn = CNN_MLP(learning_rate=0.0001)
checkpoint = None
if self.args.load_baseline_checkpoint:
checkpoint = torch.load(self.args.load_baseline_checkpoint)
cnn.load_state_dict(checkpoint['model_state_dict'])
model = cnn.to(self.device)
if self.args.use_weight_decay:
if self.args.weight_decay_all:
num_p = sum(p.numel() for p in model.parameters())
weights = np.ones(num_p) * INIT_L2
model.weight_decay = Variable(torch.FloatTensor(weights).to(
self.device),
requires_grad=True)
else:
weights = INIT_L2
model.weight_decay = Variable(torch.FloatTensor([weights]).to(
self.device),
requires_grad=True)
model.weight_decay = model.weight_decay.to(self.device)
model.train()
return model, checkpoint
def load_baseline_model(args):
"""
:param args:
:return:
"""
if args.dataset == 'cifar10':
num_classes = 10
train_loader, val_loader, test_loader = data_loaders.load_cifar10(args.batch_size, val_split=True,
augmentation=args.data_augmentation)
elif args.dataset == 'cifar100':
num_classes = 100
train_loader, val_loader, test_loader = data_loaders.load_cifar100(args.batch_size, val_split=True,
augmentation=args.data_augmentation)
elif args.dataset == 'mnist':
args.datasize, args.valsize, args.testsize = 100, 100, 100
num_train = args.datasize
if args.datasize == -1:
num_train = 50000
from data_loaders import load_mnist
train_loader, val_loader, test_loader = load_mnist(args.batch_size,
subset=[args.datasize, args.valsize, args.testsize],
num_train=num_train)
if args.model == 'resnet18':
cnn = ResNet18(num_classes=num_classes)
elif args.model == 'wideresnet':
cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10, dropRate=0.3)
checkpoint = None
if args.load_baseline_checkpoint:
checkpoint = torch.load(args.load_baseline_checkpoint)
cnn.load_state_dict(checkpoint['model_state_dict'])
model = cnn.cuda()
model.train()
return model, train_loader, val_loader, test_loader, checkpoint