pytorch 1.7.0 读取mnist数据集加载并显示图像

最新推荐文章于 2023-11-08 23:18:48 发布

Bokman

最新推荐文章于 2023-11-08 23:18:48 发布

阅读量1.1k

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签：机器学习 pytorch python

本文链接：https://blog.csdn.net/Bokman/article/details/109543304

pytorch 专栏收录该内容

3 篇文章

订阅专栏

本文展示了如何使用PyTorch加载MNIST数据集，并通过Matplotlib可视化首批次的图像。介绍了数据加载器的配置方法及图像展示的具体实现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

完整源代码：

import torch
import torchvision
import torchvision.transforms as transforms
from matplotlib import pyplot
import matplotlib.pyplot as plt
import numpy as np

batch_size = 64

train_dataset = torchvision.datasets.MNIST(root='data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='data', train=False, transform=transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)


def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    # img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    if one_channel:
        pyplot.imshow(npimg, cmap="Greys")
        pyplot.show()
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))


data_iter = iter(train_loader)
images, labels = data_iter.next()
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)

print(torch.__version__)

效果如下：