如何使用Matplotlib从PyTorch保存图形？

Posted 2021-05-07

我正在使用：

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

但是，这不会保存第一张图像。看起来像：

所以我在做什么错呢？

答案

使用matplotlib.pyplot.imshow代替matplotlib.pyplot.plot

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

这将保存matplotlib图形以仅保存可以使用matplotlib.pyplot.imsave之类的图像

matplotlib.pyplot.imsave

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

torchvision.utils.save_image