def assert_image(dataloader, index):
model = cnn.Net()
model.load_state_dict(torch.load("./modelo/mnist.pt"))
model.to(cnn.get_device())
model.eval()
image = dataloader.dataset[index][0]
image = image.to(cnn.get_device())
image = image[None]
image = image.type('torch.FloatTensor')
predictated = cnn.assert_image(model, image)
class_map = dict(map(reversed, dataloader.dataset.class_to_idx.items()))
print(class_map[predictated])
show_image(dataloader, index)
return predictated
def train_model():
dataloader = load_dataset(train=True)
model = cnn.Net()
model.to(cnn.get_device())
model.train()
model, dictionary = cnn.train_model(model=model,
dataloader=dataloader,
num_epochs=10)
print("Time spent: {:.2f}s".format(dictionary['exec_time']))
torch.save(model.state_dict(), "./modelo/mnist.pt")
save_stats(dictionary)
def test_model():
dataloader = load_dataset(train=False)
model = cnn.Net()
model.load_state_dict(torch.load("./modelo/mnist.pt"))
model.to(cnn.get_device())
model.eval()
dictionary = cnn.test_model(model=model, dataloader=dataloader)
classes = list(dataloader.dataset.class_to_idx.keys())
dictionary['classes'] = classes
dictionary = {**load_stats(), **dictionary}
save_stats(dictionary)
show_stats(dictionary)
