ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

深度学习图像分类实战:PyTorch从数据预处理到模型训练全流程

深度学习图像分类实战:PyTorch从数据预处理到模型训练全流程 1. 从零跑通一个图像分类任务我的“第一次作业”全记录事情得从一门深度学习课程说起。第一节课结束后讲师留了一项很常规的作业用任意框架实现一个图像分类模型在给定数据集上跑通训练和评估流程最终提交准确率指标和代码。听起来不复杂但对于一个此前只写过Python脚本、没碰过PyTorch和CUDA的初学者来说憋着一股劲想把它做好还是踩了不少坑。这篇内容就是这次“第一次作业”的完整复盘。我尽量按真实操作顺序来写从任务拆解、环境搭建开始到模型选型与训练调参再到中途报错和结果分析。文中涉及代码的地方我会给出可以直接复制的实现方案也会解释为什么这样写以及我在实际操作中遇到过哪些问题。无论你是刚开始接触深度学习还是准备交一份类似的课程作业/实习任务这篇记录应该都能帮你少走一段弯路。先交代一下作业的硬性要求给定一个中等规模的图像数据集大约几万张图片类别数在10个左右需要完成数据划分、预处理、模型训练、评估以及可视化结果。评价标准是测试集上的 Top-1 准确率同时会检查代码结构和实验报告的规范程度。说白了这就像完成一个最小可用的图像分类流水线从数据端到模型端到评估端全链路打通。2. 动手前的关键决策框架选择、硬件确认和任务拆解2.1 为什么我选了 PyTorch 而不是 TensorFlow在真正写代码之前我先把“用什么做”这个问题想清楚了。当时摆在面前的是 PyTorch 和 TensorFlow 二选一。两个框架都能完成图像分类任务但我最终选了 PyTorch理由有三个第一PyTorch 的调试体验更“原生”。它的动态计算图机制允许你在训练过程中随时打印张量、查看中间梯度这对第一次做深度学习项目的人来说非常友好。TensorFlow 的tf.function和静态图模式虽然性能上不差但出错时的堆栈信息对新人来说更难定位。第二PyTorch 生态里有一个非常有用的工具箱叫torchvision里面包含常用的数据集、预训练模型、图像变换函数。我们的作业数据虽然不是标准的torchvision.datasets.ImageFolder格式但也能通过自定义 Dataset 类轻松接入。这样数据加载部分的坑能少踩一大半。第三我预算有限只有一张入门级的消费级显卡。PyTorch 在显存管理和混合精度训练上的支持比较完善不需要额外配置太多东西。这个选择很务实。不同框架各有优势但作业场景下“快、稳、方便调试”才是最优先考虑的因素。如果你已经在 TensorFlow 上比较熟悉或者最终部署目标是移动端/服务端的 TensorFlow Lite那选择 TensorFlow 也完全没有问题。2.2 检查硬件再定方案开始训练之前我查看了本机环境GPUNVIDIA GeForce RTX 306012GB 显存CPUIntel i5-12400F内存16GB显存 12GB 对几万张图片的分类任务来说不算大但如果用 ResNet-50 这类模型batch size 开到 128 也能勉强跑动。我出于稳妥考虑决定先用 ResNet-18 这种参数更小的模型。为什么因为第一次作业的目标是“跑通流程”而不是“刷榜”先确保调试顺利、训练效率高后续有需要再扩大模型也来得及。检查 GPU 是否可用我用了这段代码import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU only)当时我先装的是 CPU 版本 PyTorch跑上面代码后torch.cuda.is_available()返回 False查了半天才发现安装命令少写了 CUDA 版本的标识。重新安装时用官方推荐命令加上--index-url指定 CUDA 版本的 PyTorch 仓库问题就解决了。2.3 把大任务拆成四块整份作业我没有直接写一个“大而全”的训练脚本而是拆成了四个阶段数据准备、模型构建、训练循环、评估与可视化。每个阶段先跑通一个最小 demo再糅合到最终脚本里。这种拆分方式的价值在于如果一次性把所有代码写完报错后很难定位问题。比如训练 loss 一直不下降可能是数据预处理的问题也可能是学习率设置的问题还可能是数据加载时标签对不上。分阶段验证任何一环出问题都能快速缩小范围。3. 数据加载与预处理自定义 Dataset 和标准化细节3.1 从零写一个 Dataset 类课程给的数据不是统一打包成某个常见格式的而是普通文件夹嵌套结构根目录下每个子文件夹对应一个类别图片按编号命名。正好可以用torchvision.datasets.ImageFolder来读。但问题在于有些图片可能是损坏文件还有些图片尺寸不一致直接用 ImageFolder 会在加载时出问题。我选择写一个自定义的 Dataset 类在__getitem__方法里做好异常捕获。核心代码如下from torch.utils.data import Dataset from PIL import Image import os class CustomImageDataset(Dataset): def __init__(self, root_dir, transformNone): self.root_dir root_dir self.transform transform self.classes sorted([d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]) self.class_to_idx {cls: idx for idx, cls in enumerate(self.classes)} self.images [] self.labels [] for cls in self.classes: cls_dir os.path.join(root_dir, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.images.append(os.path.join(cls_dir, fname)) self.labels.append(self.class_to_idx[cls]) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path self.images[idx] label self.labels[idx] try: img Image.open(img_path).convert(RGB) except Exception as e: print(fSkipping corrupted image: {img_path}, error: {e}) # 返回同一 batch 中的第一张图或者重新采样 return self.__getitem__((idx 1) % len(self.images)) if self.transform: img self.transform(img) return img, label最需要注意的坑就是图片损坏问题。数据集里有一小部分图片是 0KB 的空文件直接Image.open()后.load()会抛异常。我在课堂上看到不少同学遇到这个问题有的直接跳过但没重采样导致一个 batch 少了数据训练过程不稳定。我在异常处理里用了重新采样的方式如果某张图坏了就取下一张图保证返回结构完整。另一个细节是图片通道问题。有的图片是灰度图有的带透明通道.convert(RGB)可以把它们统一转成三通道避免模型输入尺寸检查报错。3.2 数据划分要分层原始数据里训练集和测试集是分开的文件夹但训练集内部没有进一步划分验证集。做深度学习时验证集对调参非常重要——你总不能每调一次超参就在测试集上跑一次那相当于作弊而且很容易过拟合。我用train_test_split做了一次分层划分保证每个类别在训练集和验证集中的比例一致from sklearn.model_selection import train_test_split train_paths [path for path in raw_images if train in path] val_paths [path for path in raw_images if val in path]如果原数据的不平衡情况比较严重光靠 random split 不够最好加stratifylabels参数。我在作业里也这么做了。3.3 数据增强和标准化图像分类的通用数据增强策略是随机水平翻转、随机裁剪、颜色抖动。这些操作本质上是“无中生有”地扩大样本数量让模型对物体的空间位置、形态变化和光照条件更鲁棒。我当时用torchvision.transforms定义了两套变换from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])讲讲这里的逻辑RandomResizedCrop(224)会随机裁剪图片的不同区域并缩放到 224x224模拟同一物体在不同尺度下的表现RandomHorizontalFlip()适合那些没有方向性的物体比如猫狗、风景、日用品ColorJitter增加光照和颜色变化的鲁棒性。验证集不需要这些随机操作所以只用Resize(256)加CenterCrop(224)。标准化用的均值和标准差直接选了 ImageNet 的统计值。因为 ResNet 系列预训练模型的统计就是在这些数值下算出来的如果你后面打算加载预训练权重这一步必须跟预训练时的设置一致否则模型的输入分布对不上。4. 模型选型预训练 ResNet-18 还是从零训练4.1 预训练模型的优势第一次做图像分类最容易纠结的就是“模型自己从零训练好还是用预训练权重微调好”。我当时的想法比较偏理想化觉得自己写一个 CNN 从零训练能学到更多。但实际跑了几轮后发现效果非常一般因为几万张图片对一个从零开始的深度网络来说根本不够模型很容易过拟合到训练集上验证集准确率徘徊在 70% 上下。后来改用 torchvision 里预训练好的 ResNet-18只替换最后一层全连接在相同条件下跑验证集准确率直接跳到了 90% 以上。差距非常大原因很简单预训练模型已经在 ImageNet 那 1000 类的海量图片上学到了大量通用特征比如边缘、纹理、形状结构你只需要在它的基础上微调一个小分类头就能把它学到的特征迁移到你的任务上。4.2 替换分类头的代码ResNet-18 的最后一层是一个全连接层输入维度是 512输出维度是 1000。我需要把它替换成输出类别数的全连接层import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) num_features model.fc.in_features num_classes 10 model.fc nn.Linear(num_features, num_classes)这里需要注意weights参数。新版本 torchvision 要求显式指定 weights否则可能只会加载随机初始化的模型起不到迁移学习的效果。很多人以为即使不冻结特征提取层只换最后一层就行。事实上默认情况下整个模型的所有参数都会参与梯度更新这被称为全量微调fine-tune。如果你希望训练速度更快、显存占用更少可以先把特征提取层冻结只训练最后的分类头for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True我在第一次作业中选择了全量微调因为数据集规模还可以全量微调上限更高。但刚开始调试时为了快速验证代码是否跑得通先用过冻结特征层的方案训练速度提升了将近一倍。建议你也这样先冻结跑通再解冻跑全量。4.3 为什么不选用更大的模型我周围有同学一上来就选了 ResNet-50甚至是 EfficientNet-B4。模型参数多、表达能力强但对数据和显存的要求也更高。RTX 3060 的 12GB 显存跑 ResNet-50 的 batch size 32 也还可以但训练时间会明显增加。作为第一次作业我更看重“快速迭代”的体验所以先用 ResNet-18 拿到一个可靠的基线时间充裕再尝试更复杂的模型。这里也想提醒一点刷高分是目标之一但理解训练流程才是真正重要的收获。如果第一次就堆一个超大模型复现起来全是资源问题反而会掩盖学习重点。5. 训练循环中的三个关键点损失函数、优化器、学习率5.1 用交叉熵损失别走弯路分类任务最常用的损失函数是交叉熵损失。PyTorch 里nn.CrossEntropyLoss可以直接结合最后一层线性输出不需要手动做 Softmax因为在损失函数内部已经把 Softmax 和对数损失合并在一起了。我之前看到有教程先对输出做F.softmax再传给它虽然结果相差不大但本质多此一举数值稳定性还可能受到影响。直接这样写就够了criterion nn.CrossEntropyLoss()5.2 优化器选择Adam 适合起步SGD 精调更稳优化器我用了 Adam理由是它对学习率没那么敏感能自动调节每一个参数的学习步长。这对第一次写训练循环的人来说能减少调参白头发。Adam 的默认学习率是1e-3我在作业里直接使用前期效果就不错。但我后来也发现Adam 收敛到后期往往不如带动量的 SGD 来得“扎实”。这是因为 Adam 会对历史梯度做指数移动平均后期更新步长受二阶矩影响可能出现在最优解附近来回震荡的情况。所以我做了两段式训练前 10 个 epoch 用 Adam学习率1e-3把模型从一个随机初始化状态快速拉到一个合理的局部最优附近然后切到 SGDmomentum0.9学习率降到1e-2或1e-3进一步精调。切换优化器后验证集准确率又涨了大概 1 到 2 个点效果还挺明显。当然这不是标准答案不同任务表现可能不同但可以作为调优路径的一个参考。5.3 学习率调度不能少训练过程中如果一直保持同一个学习率很容易在 loss 平台期来回震荡。我用的是余弦退火调度器from torch.optim.lr_scheduler import CosineAnnealingLR scheduler CosineAnnealingLR(optimizer, T_max20, eta_min1e-6)它的含义是在一个周期内学习率从初始值按余弦曲线下降到eta_min。这样训练的后期步长自动变小可以精细地逼近最优解。换成这个调度器后最终准确率又稳了一截而且训练曲线看起来顺滑得多。5.4 训练循环的标准模板下面这份代码是我这次作业训练循环的核心骨架每一段我都在旁边加了注释def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc一个很容易被忽略的细节是optimizer.zero_grad()。PyTorch 的梯度是默认累加的如果不清零多个 batch 的梯度会叠在一起loss 会表现得异常训练也无法收敛。这个坑我踩过一次训练 loss 莫名其妙一直往上走后来发现就是忘了在每个 batch 开头调用zero_grad()。还有个细节是model.train()和model.eval()的切换。模型里如果有 BatchNorm、Dropout它们的训练行为与推理行为完全不同。BatchNorm 在训练时用当前 batch 的统计量在推理时用全局移动平均统计量Dropout 在训练时随机丢弃节点在推理时保持全通。所以验证阶段必须切换到 eval 模式否则验证指标会很离谱。6. 训练实测与调优过程从 loss 不降到最终收敛6.1 第一次跑loss 不降反升我清清楚楚记得第一次完整跑完一个 epoch 的心情训练 loss 不但没下降反而从 2.3 涨到了 2.8。当时第一反应是代码写错了于是回到 DataLoader 去检查标签和数据是否对得上。排查思路是这样的先取一个 batch 的数据把图片和标签都打印出来人工看了一眼发现标签和图片内容是匹配的。然后又检查了模型输出维度类别数也对。最后去查学习率才意识到问题可能出在数据预处理的标准化上——如果训练集和验证集都用了同样标准化算出来的分布倒还好但我训练集用的是RandomResizedCrop验证集用的是CenterCrop两者对图片的分布影响其实很小可训练 loss 不降从经验上看大概率是学习率太高或者优化器状态有问题。最终定位到我把学习率设成了5e-3对 Adam 来说偏高梯度更新幅度过大导致 loss 发散。降到1e-3后loss 立刻恢复了正常下降趋势。这再次说明新手应该先尽量遵循默认超参再在此基础上小幅调整不要一上来就追求“激进”。6.2 正常训练但是验证集准确率卡在 82%调整完学习率后训练变得很顺利大约 8 个 epoch 后验证集准确率到了 82%然后就像被“粘住”了一样怎么都上不去。这个阶段我做了几件事把训练集增强加得更强增加了RandomRotation(10)又加了RandomGrayscale(p0.1)把模型从冻结特征层改为全量微调把 batch size 从 64 减到 32并同步调低了学习率。这几项调整合在一起准确率才慢慢从 82% 爬升到了 89%。回过头看冻结特征层时训练的是随机初始化的分类头它需要尽快拟合任务但由于特征层固定成 ImageNet 的通用特征对当前数据集不一定最优。解冻后特征层会朝任务数据的方向适应自然能带来提升。6.3 过拟合迹象与早停训练到后面我还观察到训练准确率已经到了 97%但验证准确率始终停在 90% 上下这是典型的过拟合信号。解决办法我用了两种一是把 Dropout 加到分类头里虽然 ResNet 本身没有 Dropout但可以在 fc 层前插入二是早停机制。早停代码很简单但很实用best_val_acc 0.0 patience 5 bad_epochs 0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) bad_epochs 0 else: bad_epochs 1 if bad_epochs patience: print(fEarly stopping at epoch {epoch}) break保存“最佳模型”的时机一定是在验证集准确率最高的时候而不是最后一个 epoch。如果只是简单保存最后一个 epoch 的权重可能已经过拟合实际部署效果会打折。6.4 最终测试集结果用保存的最佳权重跑测试集最终 Top-1 准确率是 91.6%。跟班级里的高分比不算顶尖有人做到 94%但整个过程中把数据、训练、评估全链路都亲手跑通了一遍收获是实在的。作业达标之外还积累了一套可以复用的代码模板。7. 训练过程中的常见报错与解决方法写代码哪有不出错的第一次作业里我至少遇到了五类乱七八糟的报错这里挑几个典型列出你们以后大概率也能碰到。7.1 CUDA out of memory错误信息RuntimeError: CUDA out of memory.我第一次遇到时很慌后来发现是另一个进程占了一部分显存。用nvidia-smi查看后把无关进程清掉问题就解决了。如果确实显存不够几个立竿见影的技巧减小 batch size比如 64 降到 32减小图片尺寸比如 224 降到 160但这样可能影响准确率用混合精度训练torch.cuda.amp能在不明显掉精度的情况下降低显存占用。混合精度训练是显存不足时的“解药”级别操作在 PyTorch 里只需要三行代码scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()我在训练后期开了混合精度显存占用降了将近 40%训练速度还快了不少。7.2 Expected all tensors to be on the same device错误信息通常长这样Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!原因很简单模型和输入数据有的在 GPU 上有的在 CPU 上。解决办法是统一放到同一个设备device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) images, labels images.to(device), labels.to(device)7.3 图像尺寸不匹配加载预训练 ResNet 后如果你的输入图片不是 224x224最后一层 fc 层的维度就对不上。一定记住RandomResizedCrop的作用它一方面做数据增强另一方面也负责把尺寸统一。7.4 断点续训和训练中断有一次训练跑到了第 12 个 epoch电脑因断电中断了。辛辛苦苦跑了一晚上的结果说没就没令人崩溃。后来我写了检查点保存逻辑每完成一个 epoch 就保存一份checkpoint.pth内容包括模型权重、优化器状态、调度器状态和当前 epochtorch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), best_val_acc: best_val_acc, }, fcheckpoint_epoch_{epoch}.pth)恢复训练时加载这些状态再继续即可。这个习惯救了我第二次工作室同学有一次在调度任务到期后训练中断了靠检查点马上续上少跑了一整轮。8. 评估与可视化不只是打印一个准确率数字作业要求里有一项是“可视化评估结果”很多同学只打印了一个测试集准确率就交给老师了。我额外做了几件事也强烈建议你这样做因为从评估可视化里能看到比单一指标多得多的信息。8.1 混淆矩阵我生成了测试集上的混淆矩阵并用 seaborn 热力图展示。混淆矩阵能直观告诉你模型把哪些类别互相搞混了。我们的数据里“鸡”和“鸭”互认错的情况最多我一看混淆矩阵就明白了。后来专门查看了这些误判样本发现不少图片的拍摄角度都很刁钻背景杂乱人眼都容易搞错模型有这种失误也情有可原。8.2 按类别统计准确率除了整体准确率我还按类别统计了每个类别的准确率、召回率、F1。整体 91.6% 的准确率很亮眼但其中某个类别准确率只有 80%如果只看整体数字发现不了这回事。分拆统计后才能知道后续优化应该优先处理哪些类别。8.3 可视化训练曲线把每个 epoch 的 train_loss、val_loss、train_acc、val_acc 画在一张图上是判断过拟合、学习率是否合理的直接手段。我自己的训练曲线val_loss 从第 12 个 epoch 开始不再下降但 train_loss 还在降这就是一个再明显不过的过拟合信号。9. 复盘一次作业下来我沉淀了哪些可复用的套路做完这次作业我自己总结了一套可迁移到后续项目中的“图像分类任务标准流程”。拿到数据先做 EDA统计类别数、每类样本数、图片尺寸分布、损坏文件数量这一步能指导后续预处理方法先建立完整且能跑的 baseline用相对小的模型、相对快的 epoch 数验证整个流程没有问题后再做优化优化优先顺序从“数据处理→采样策略→数据增强→网络结构→超参”每一步只改变一个变量便于归因训练过程中每隔几轮保存检查点不要关掉终端就什么都不留评估阶段不只看准确率混淆矩阵、类别指标、训练曲线三件套一个都不要少。这篇文章基本把我第一次做图像分类作业的过程完完整整还原了一遍。如果非要浓缩成一句话我会说第一次做项目最重要的不是一开始就追求高分而是把“数据处理→模型训练→评估调优”的闭环跑通从中积累一套属于自己的稳定模板。有了这套模板后面换数据集、换模型、换任务方向时核心逻辑都不需要大改。如果你现在正卡在某个报错上或者训练准确率上不去试着按我文章里的排查顺序过一遍先看数据对不对再看设备和 dtype 是否统一然后是学习率是否合理最后才是网络结构。多跑几个 epoch把训练曲线画出来很多问题一眼就能看出来。第一次作业踩坑不丢人能从坑里爬出来并且能说清自己是怎么爬出来的这本身就是最大的收获。
返回列表