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