def get_train_val_test_loader(self, args): split_to_data = {} for split in ['train', 'val', 'test']: data_prep = DataPreparation(args.dataset, args.data_path) dataset, _ = data_prep.get_dataset_and_loader(split, args.pretrained_model, batch_size=args.batch_size, num_workers=args.num_workers) split_to_data[split] = dataset return split_to_data
def get_model(): old_args = sys.argv[:] sys.argv = old_args[:1] # Get default arguments args = arg_parser.get_args() sys.argv = old_args args.model = "gve" args.dataset = "cub" args.pretrained_model = "vgg16" args.num_epochs = 1 args.batch_size = 1 # set to train because we need gradients for Grad-CAM args.train = True args.eval_ckpt = "data/vgg-ic-gve-best-ckpt.pth" args.ic_ckpt = "data/cub/image_classifier_ckpt.pth" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Data preparation print("Preparing Data ...") split = get_split_str(args.train, bool(args.eval_ckpt), args.dataset) split = "test" data_prep = DataPreparation(args.dataset, args.data_path) dataset, data_loader = data_prep.get_dataset_and_loader( split, args.pretrained_model, batch_size=args.batch_size, num_workers=args.num_workers, ) # Load VGE model print("Loading Model ...") ml = ModelLoader(args, dataset, device) model = getattr(ml, args.model)() print(model, "\n") print("Loading Model Weights ...") evaluation_state_dict = torch.load(args.eval_ckpt, map_location="cpu") model_dict = model.state_dict(full_dict=True) model_dict.update(evaluation_state_dict) model.load_state_dict(model_dict) # Disable dropout and batch normalization model.eval() model.has_vision_model = False vgg_feat_layers = ( model.image_classifier.vision_model.pretrained_model.features) vgg_class_layers = None visual = np.zeros((224, 224)) trainer_creator = getattr(TrainerLoader, args.model) trainer = trainer_creator(args, model, dataset, data_loader, logger=None, device=device) return model, trainer, dataset, vgg_feat_layers
args.dataset = 'cub' args.pretrained_model = 'vgg16' args.num_epochs = 1 args.batch_size = 1 # set to train because we need gradients for Grad-CAM args.train = True args.eval_ckpt = 'data/vgg-ic-gve-best-ckpt.pth' args.ic_ckpt = 'data/cub/image_classifier_ckpt.pth' device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Data preparation print("Preparing Data ...") split = get_split_str(args.train, bool(args.eval_ckpt), args.dataset) split = 'test' data_prep = DataPreparation(args.dataset, args.data_path) dataset, data_loader = data_prep.get_dataset_and_loader( split, args.pretrained_model, batch_size=args.batch_size, num_workers=args.num_workers) # Load VGE model print("Loading Model ...") ml = ModelLoader(args, dataset, device) model = getattr(ml, args.model)() print(model, '\n') print("Loading Model Weights ...") evaluation_state_dict = torch.load(args.eval_ckpt, map_location='cpu') model_dict = model.state_dict(full_dict=True) model_dict.update(evaluation_state_dict)