def load_model(): ''' Loads a pre-trained model and settings used to generate it. ''' try: with open(f'{PATH_TO_MODEL}.json', 'r') as json_file: settings = json.load(json_file) model = NeuralNet(len(settings['all_labels'])) model.load_state_dict( torch.load(f'{PATH_TO_MODEL}.pth', map_location=torch.device('cpu'))) except: print('Could not locate a trained model.') sys.exit() model.eval() return model, settings
data_x, data_y = valid_x, valid_y test_dataset = Dataset(root_dir, data_x, data_y, transforms=transform) test_generator = torch.utils.data.DataLoader(test_dataset, **params) print("Loaded dataloaders...") criterion = torch.nn.CrossEntropyLoss() model = NeuralNet(0.001, criterion, 64, 2) state_dict = torch.load(model_name) model.load_state_dict(state_dict) for parameter in model.parameters(): parameter.requires_grad = False if (use_cuda): model.cuda() model.eval() summary(model, (1, 64, 64)) print("Loaded model...") preds = [] labels = [] for local_batch, local_labels in tqdm(test_generator): labels.extend(local_labels.numpy().tolist()) local_batch, local_labels = local_batch.to(device), local_labels.to(device) output = model.forward(local_batch) output = model.softmax(output) output = torch.max(output, 1)[1] preds.extend(output.cpu().numpy().tolist())