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())
