How do I save a figure from PyTorch using Matplotlib?

Question

I'm using:

trainset = datasets.MNIST('saved/train', download=True,
                          train=True, transform=transform)

valset = datasets.MNIST('saved/test', download=True,
                        train=False, transform=transform)

trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True)
valloader = torch.utils.data.DataLoader(valset, batch_size=64, shuffle=True)


dataiter = iter(trainloader)
images, labels = dataiter.next()

fig = plt.figure()
print(images.shape)
print(labels.shape)

plt.plot(images[0].numpy().squeeze())

fig.savefig('figs/first.png')

However, this does not save the first image. It looks like:

So what am I doing wrong?

Answer 1

Use matplotlib.pyplot.imshow instead of matplotlib.pyplot.plot

plt.imshow(images[0].numpy().squeeze())
fig.savefig('first_fig.png')

This will save matplotlib figure to save only image you can use matplotlib.pyplot.imsave like

plt.imsave('first_imsave.png',images[0].numpy().squeeze())

or use torchvision.utils.save_image

utils.save_image(images[0],'first_utils.png')