utils.py

import torch
import matplotlib.pyplot as plt
import torchvision

mnist_raw = torchvision.datasets.MNIST('./files/', train=False, download=True,
 transform=torchvision.transforms.ToTensor())


def test_single_image(network, example_number):
 with torch.no_grad():
 fig = plt.figure()
 ax = fig.add_subplot(1,1,1)
 image, label = mnist_raw[example_number]
 ax.imshow(image[0], cmap='gray', interpolation='none')
 ax.set_xticks([])
 ax.set_yticks([])
 image = image.reshape(-1, 28*28)
 output = network(image)
 ax.set_title("Korrekt: {}, Vorhersage: {}".format(label, torch.max(output.data,1).indices[0]))
 plt.show()


def draw_single_digit(number):
 fig = plt.figure()
 ax = fig.add_subplot(1,1,1)
 image, label = mnist_raw[number]
 ax.imshow(image[0], cmap='gray', interpolation='none')
 ax.set_title("Label: {}".format(label))
 ax.set_xticks([])
 ax.set_yticks([])