首页 / 资讯中心 / 文章详情

零基础PyTorch加载MNIST数据集:从环境到DataLoader实战

零基础PyTorch加载MNIST数据集:从环境到DataLoader实战 ★ FEATURED ARTICLE
第二天我终于把昨天装好的 PyTorch 环境打开了。说实话在真正碰到数据集之前我一直觉得“深度学习”这四个字离我很远直到亲手把 MNIST 跑起来才发现它其实没那么神秘。这篇是“零基础使用 PyTorch 跑 MNIST”系列的第二篇内容就一件事怎么把 MNIST 数据集装进程序里让它变成模型能读、能算、能练手的数据。适合刚装完环境、准备迈出第一步的人也适合那些已经跑过代码但一直没搞明白数据集加载过程的朋友。1. 先搞清楚MNIST 数据集到底“长什么样”1.1 像素、灰度图与数字标签MNIST 是手写数字识别领域最经典的数据集里面全是 0 到 9 的手写数字图片。每张图片的尺寸是 28×28 像素而且是单通道的灰度图不是彩色图。你可以把它想象成一张 28 行、28 列的小方格纸每个格子里存一个 0 到 255 之间的整数数值越大代表这个格子越亮0 就是纯黑255 就是纯白中间的数字是不同程度的灰色。整个数据集一共包含 60000 张训练图片和 10000 张测试图片每张图片都配了一个标签标签就是这张图片对应的真实数字比如一张写着“7”的图片它的标签就是整数 7。这里有个零基础很容易忽略的点图片在电脑里不是一个整体它只是一个二维数组加一个标签。PyTorch 加载数据集时做的事情本质上就是把一堆这样的“二维数组 标签”组合在一起打包成一个方便操作的对象。你不需要了解底层文件格式的每个字节但必须理解“图片 像素矩阵、标签 数字”这个基本结构否则后面查 bug 时很容易迷失方向。1.2 为什么零基础第一个项目都选它很多人好奇MNIST 都快被跑烂了为什么所有入门教程还选它我自己的体会是它有三个无可替代的优点。第一数据量适中60000 张图不算大一张图才 28×28单张只要几百字节整套数据集下载下来也不过几十 MB普通电脑完全扛得住。第二任务本身足够简单手写数字识别不需要复杂的网络结构用一个最简单的全连接网络就能达到很高的准确率非常适合用来理解“训练、验证、测试”这一套完整流程。第三PyTorch 官方工具库 torchvision 内置了 MNIST 的下载和加载接口你不用自己去网上找数据、自己写解析代码几行代码就能完成这能帮你把精力集中在“学习框架本身”而不是被数据预处理劝退。不过我要提醒一句正因为 MNIST 太成熟了网上各种代码满天飞新手很容易照抄完就觉得自己会了。我建议你在跑通之后一定要亲手查看一下每个数据集对象的类型、长度、元素结构把返回的图片张量形状和标签值打印出来看一遍。这一步比跑通代码本身更重要。2. 环境确认与核心依赖torchvision 才是加载数据的“主角”2.1 动手前先花两分钟确认环境第一天的环境搭建如果顺利现在你的电脑里应该有 PyTorch 和 torchvision。加载 MNIST 这件事PyTorch 本体只负责张量运算真正负责下载和读取数据的是 torchvision它是 PyTorch 官方配套的视觉工具库。打开终端或编辑器先跑一段小代码确认环境没问题import torch import torchvision print(torch.__version__) print(torchvision.__version__)能正常打印出两个版本号就说明基础环境 OK。如果运行时报ModuleNotFoundError: No module named torchvision说明你只装了 torch 没装 torchvision回到安装环节把 torchvision 补上。这里特别提醒一点torch 和 torchvision 之间存在版本对应关系如果你是用 pip 分别安装的最好安装时指定兼容的版本号否则会碰到莫名其妙的报错比如后面加载数据时提示某个属性不存在其实不是你的代码问题而是版本不匹配。2.2 torchvision.datasets.MNIST 做了什么torchvision.datasets.MNIST 是官方为你封装好的数据集类。它内部干了几件重要的事检查指定目录下有没有 MNIST 数据没有的话就自动去官方地址下载有的话就直接读取再根据你传入的 transform 参数对图片做预处理最后把数据整理成 PyTorch 能直接操作的数据集对象。这里面的“检查目录”和“自动下载”是新手最容易踩坑的点因为它的逻辑是“文件在就直接用文件不在就下载”如果下载因为网络原因失败了一半留下的不完整文件会导致后续每次运行都报错。一个更值得留意的点是MNIST 官方数据存储的格式和常见图片格式比如 JPG、PNG完全不同。它使用了一种名为 IDX 的二进制格式分为图片文件和标签文件。torchvision 的价值就在于把这些底层细节全部屏蔽掉了你只需要通过参数告诉它数据存在哪、是否需要下载它就能把图片和标签一起读出来。如果让你自己写解析代码虽然也能写但至少要多花一两个小时而且很容易在小细节上出错。我的建议是刚开始不要重复造轮子老老实实用官方接口等后面理解深了再去看源码都来得及。2.3 下载卡住怎么办手动方案与镜像策略自动下载是默认行为但国内网络环境下直接从官方地址下载经常卡住或超时。我第一天就被这个坑折磨过看起来好像在下载进度条半天不动最后直接报超时错误。如果你也碰到这种情况可以采取手动下载策略先从一些国内开源镜像站点找到 MNIST 的四个文件训练图片、训练标签、测试图片、测试标签下载后放到项目目录下然后在代码里把download参数设置为False这样 torchvision 就会直接读取本地文件不再触发联网下载。更稳妥的做法是在使用datasets.MNIST时手动指定root参数为你放置好文件的目录。这里有个细节torchvision 会自动识别目录里的原始文件并生成一个processed子目录存放转换后的二进制缓存。如果你之后改了代码里某些参数它可能会重新处理第一次运行会稍微慢一点但第二次就快很多。注意千万不要手动去改 processed 目录里的文件一旦缓存损坏直接删除整个 MNIST 目录让它重新下载或重新处理就行。3. 第一段代码创建数据集对象并读懂返回结果3.1 参数逐个拆解现在写第一段真正的加载代码。建议新建一个项目文件夹在文件夹里建一个 Python 脚本内容和下面这段保持一致import torchvision import torchvision.transforms as transforms # 定义数据预处理把 PIL 图片转成 Tensor并把像素值归一化到 [0, 1] transform transforms.Compose([ transforms.ToTensor(), ]) # 加载训练集 train_dataset torchvision.datasets.MNIST( root./data, # 数据存放目录 trainTrue, # True 表示加载训练集False 表示加载测试集 transformtransform, # 预处理操作 downloadTrue # 如果本地没有数据就自动下载 ) # 加载测试集 test_dataset torchvision.datasets.MNIST( root./data, trainFalse, transformtransform, downloadTrue ) print(训练集大小:, len(train_dataset)) print(测试集大小:, len(test_dataset))运行后你应该看到类似训练集大小: 60000、测试集大小: 10000的输出。很多教程会直接让你复制这段代码然后继续但我强烈建议你花两分钟盯一下参数因为后续所有数据处理都建立在这几个参数上。root是数据根目录如果目录不存在torchvision 会自动创建。train参数决定加载哪一部分数据这是新手经常混乱的地方第一次运行会把 60000 张训练图和 10000 张测试图全都下载下来train 只是决定当前这个对象读取哪部分。transform是数据预处理的入口目前只有一个ToTensor()。downloadTrue的意思是允许自动下载但如果之前已经下载好了它不会重复下载而是直接复用。3.2 用下标访问数据集亲手摸一摸数据和标签MNIST 数据集对象支持像列表一样用下标访问。每次访问会返回一个元组里面是两部分图片数据和对应的标签。你可以用下面这段代码把第一条数据拿出来看看image, label train_dataset[0] print(图片数据类型:, type(image)) print(图片张量形状:, image.shape) print(图片像素值范围:, image.min().item(), -, image.max().item()) print(标签:, label, 标签类型:, type(label))你会看到类似这样的输出torch.Size([1, 28, 28])标签是一个整数。这里的[1, 28, 28]中的 1 是通道数因为 MNIST 是灰度图所以是 1 个通道如果你处理彩色图这个位置会是 3。很多人第一次看到这个形状会疑惑为什么不是[28, 28]这是 PyTorch 的约定张量的第一个维度代表通道数后面两个维度才是高和宽。后面训练卷积网络时这种格式的优势就体现出来了所以从第一天开始就要习惯它。image的类型是torch.Tensor这是经过ToTensor()转换后的结果。你还可以用它来验证数据是否正确比如打印image[0][14][14]取到的是第 14 行第 14 列那个像素的值。我实际操作时喜欢把第一个样本的 28×28 像素用文本方式打印出来看一遍虽然很费眼睛但能立刻建立起“图片原来就是这样一组数字”的直观感受强烈推荐零基础的朋友做一次这个小实验。3.3 transformtransforms.ToTensor() 那点事transforms.ToTensor()是加载过程中最关键的预处理步骤它做了两件事一是把 PIL 图像或 NumPy 数组转换成 PyTorch 的张量二是把像素值从 0 到 255 的整数缩放到 0 到 1 的浮点数。你可能会想为什么非要缩放直接用 0 到 255 的整数不行吗这里有个深度学习的基本常识神经网络训练过程中数值尺度太大会让梯度计算变得不稳定让损失函数来回震荡进而导致模型很难收敛。而归一化到 0 到 1 之后输入数据的分布更平稳训练更容易稳定下来。就像你去健身房举杠铃如果第一次就上 100 公斤大概率举不起来还容易受伤把重量调到合适范围才能保证动作标准、持续进步。同理如果你后面想用均值 0、标准差 1 的标准正态分布来进一步归一化还可以继续在 Compose 里加transforms.Normalize((0.1307,), (0.3081,))这两个数值正是 MNIST 数据集所有像素的全局均值和标准差但在入门阶段只用ToTensor()就够用了。要注意的是ToTensor()之后张量的值虽然是 0 到 1 的浮点数但它只是在数值范围上做了缩放图片本身的明暗对比、内容结构完全没有改变。你后面如果要把图片可视化出来就需要把像素值再乘回 255或者告诉绘图库数据范围是 0 到 1否则画出来是一团黑色或白色。4. DataLoader把零散图片变成“模型能一口吃下的小饭团”4.1 四个关键参数怎么调有了数据集对象之后下一步不是立刻训练而是把它丢进 DataLoader。DataLoader 是 PyTorch 提供的数据加载器它会把数据集里的样本按一定规则打包成一个个批次并在训练过程中按顺序输出。同样的先上代码from torch.utils.data import DataLoader train_loader DataLoader( datasettrain_dataset, batch_size64, shuffleTrue, num_workers0 ) test_loader DataLoader( datasettest_dataset, batch_size64, shuffleFalse, num_workers0 ) print(训练批次数量:, len(train_loader))batch_size表示每批包含多少张图片。我用的是 64这是入门阶段非常常见的取值。假设一次把全部 60000 张图片都塞给模型显存瞬间爆掉不说训练也没法高效进行。所以让模型每次只看 64 张计算一次梯度更新一次参数然后再看下一批。len(train_loader)会告诉你这批数据集被分成了多少份64 的话大概是 938 份60000 除以 64不满一份的直接丢弃。shuffleTrue表示每个 epoch 开始时重新打乱数据顺序。为什么要打乱想象你背单词如果每次都按同一个顺序背你可能会记住单词的位置而不是单词本身考试换个顺序就蒙了。模型也是这样如果不打乱它对特定顺序会产生依赖泛化能力会变差。测试集就绝对不能打乱因为测试只是评估模型效果不更新参数顺序无所谓但保持原始顺序能保证评估结果可复现。num_workers表示用几个子进程来加载数据。零基础阶段我建议直接填 0也就是在主进程里加载简单而且不容易出问题。在 Windows 上尤其要注意如果设成大于 0 的数字且在 Jupyter Notebook 里运行很容易触发子进程递归启动的报错搞不好直接死机这个坑我帮你们踩过了。4.2 批量取出并可视化第一小批样本光看数字肯定不够直观最好把图片画出来看看。可视化不仅能验证数据加载是否正常还能让你对“数据集长什么样”有直接印象。下面这段代码可以画出第一个批次里的前 10 张图import matplotlib.pyplot as plt # 从迭代器里取一批数据 data_iter iter(train_loader) images, labels next(data_iter) print(批次图片形状:, images.shape) print(批次标签形状:, labels.shape) # 画前 10 张 plt.figure(figsize(10, 4)) for i in range(10): plt.subplot(2, 5, i 1) plt.imshow(images[i].squeeze(), cmapgray) plt.title(flabel: {labels[i].item()}) plt.axis(off) plt.show()images.shape会输出torch.Size([64, 1, 28, 28])这是一个四维张量批次大小 64、通道数 1、高 28、宽 28。labels.shape是[64]。这里的squeeze()很重要因为images[i]的形状是[1, 28, 28]imshow不希望看到那个多余的通道维度squeeze()会把它去掉变成[28, 28]。如果你想看其中一张图的像素值分布可以用下面这段代码single_image images[0].squeeze().numpy() plt.imshow(single_image, cmapgray, vmin0, vmax255) plt.colorbar() plt.show()这里vmin和vmax参数非常关键因为张量值在 0 到 1 之间如果你不设置这两个参数matplotlib 可能会自动调整色彩映射范围导致图片看起来全是黑的或者对比度怪怪的。设置成 0 到 255或者干脆先不设置而是显示原始 0 到 1 范围图片都会正常显示。我第一次就是没设这个参数画出来一片漆黑还以为数据加载错了查了半天才发现是可视化范围的问题。4.3 常见的维度报错与排查零基础阶段报错最多的地方就在数据维度和类型上。我把常见的几类问题整理一下你碰到时可以对照着排查。第一个是“图片形状不对”的错。比如你预期输入是四维张量但实际的 shape 却是[28, 28]或[1, 28, 28]。这种情况通常是你在某个位置多写或少写了维度比如在自定义模型里用了view(-1, 28*28)把三维张量压成二维时数字没算对。MNIST 每张图是 784 个像素也就是 28×28所以展平后应该是[-1, 784]如果你写成了[-1, 28]模型根本接收不了。第二个是设备不匹配的问题。数据默认在 CPU 上模型如果被挪到了 GPU那么数据也必须挪过去需要调用images images.to(device)。新手经常会忘记这一步报错提示会明确告诉你 expected device cuda but got device cpu看到这个提示不要慌去检查数据有没有to(device)就行。第三个是数据类型不匹配的问题。模型里的权重一般是torch.float32如果输入数据是torch.uint8或torch.int64某些操作就会报错。用ToTensor()之后图片已经变成浮点类型了所以这类问题在 MNIST 里不常见但当你自己往流程里加了别的数据处理逻辑时就要时刻留意数据类型。第四个是迭代器用尽的问题。iter(train_loader)创建的迭代器在循环完一遍之后会耗尽如果再对它调用next()就会报StopIteration。很多人在调试或写自定义训练循环时会踩到这个坑解决办法是每次需要取数据时重新创建一个data_iter iter(train_loader)。另外enumerate(train_loader)返回的 batch 序号是在当前 epoch 内的序号epoch 结束就从头再来这个逻辑要理清。5. 常见问题速查表与几点实操心得5.1 问题排查速查表现象可能原因解决办法下载时网络超时或卡住直连官方数据源不稳定从国内镜像站点手动下载四个文件放到./data目录downloadFalse提示Dataset MNIST not foundroot 路径不对或文件放错位置检查root路径确认文件在对应raw子目录下实在不行删掉目录重新下载每次运行都重复下载目录下没有完整文件或 download 参数误设为 True 且 root 变了检查文件完整性统一 root 路径删除不完整的损坏文件数据加载后图片全黑imshow色彩范围没设置对添加vmin0, vmax1或把像素乘回 255 范围Windows 下 num_workers 大于 0 导致卡死子进程递归问题把num_workers设为 0报错Expected tensor to have 4 dimensions把三维或二维张量直接送入模型检查是否缺少 batch 维度必要时使用unsqueeze(0)添加维度报错Expected device cuda but got device cpu数据没做设备迁移调用images.to(device)和labels.to(device)打印数据集长度时报内存相关错误数据缓存异常或磁盘空间不足检查磁盘空间删除processed缓存目录重新生成同一个 batch 的图片都一模一样没有正确打乱数据或读取错误确认训练 loader 的shuffleTrue并打印多个样本核对5.2 零基础最容易忽略的细节第一个细节数据集的加载不是一次性把 60000 张图片全部堆在内存里。DataLoader 的机制是每取一个批次才加载对应图片执行完操作后这批数据的引用就释放了。如果你的机器内存不够可以考虑把batch_size调小比如从 64 调到 32 甚至 16。但注意batch_size 太小会导致训练不稳定太大则一次计算的样本太多、内存吃紧这个平衡需要根据你的实际硬件来调没有绝对标准。第二个细节shuffleTrue只作用于训练集。如果训练时不小心把它用在测试集上唯一直接后果是每次测试顺序不同不会影响准确率计算但会让调试过程变得混乱比如你想对比两次测试的输出就发现无法对齐。所以养成习惯训练用 shuffle测试不用。第三个细节数据增强要谨慎加。MNIST 本身已经很规范入门阶段不需要任何数据增强。有些教程为了演示会加上随机旋转、随机裁剪之类操作但如果你追求的是一步一步理解框架建议先不加等模型能正常训练之后再对比增加这些操作前后的效果差异。这样你能直观感受到数据增强带来的影响而不是糊里糊涂地把所有代码堆上去。第四个细节是路径问题。如果你的项目换了目录之前的下载数据不会自动跟着搬过去重新跑代码时它检测不到数据又会触发下载。所以我建议把root./data固定下来整个项目周期内都使用同一个数据目录别这个脚本用data/mnist那个脚本用./MNIST_data到时候白白浪费流量。还有一点如果数据目录在代码中被设为相对路径运行脚本时的工作目录不同也会导致找不到数据这个坑也常遇到可以用绝对路径或确保每次都从项目根目录启动脚本来规避。最后聊聊我个人的体会。加载 MNIST 是整个深度学习入门流程里最简单、但也是最重要的一环。说它简单是因为代码就那么几行几分钟就能跑通说它重要是因为你后面写的训练循环、评估脚本、模型前向传播全都建立在对“数据集对象、DataLoader、批次张量”这几个概念的正确理解上。我在第一次跑通时其实只花了十几分钟但后面反复研究 shape 和 DataLoader 的取数逻辑花了一下午。我不建议你用“跑通就行”的心态糊弄过去而是建议你把每一行打印结果都看一遍甚至动手把dataset[0]、dataset[0][0]、dataset[0][1]分别打印出来搞清楚每一项到底是什么。还有一个对零基础特别友好的排查思路遇到报错时先把报错信息带着的 shape 信息全部抄到纸上然后对照“模型期望输入的形状”和“当前数据的形状”一条一条核对。90% 的数据相关报错都能通过这种方式解决。等你熟练了这个过程接下来的训练、评估、调参就只是在这个基础上继续添砖加瓦而已。第 2 天到这里数据已经把路子铺好了下一篇可以真正开始写模型。
阅读完成 · 觉得有帮助?
咨询建站