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

从零构建AI图像分类器:PyTorch实战Doompop项目全流程

从零构建AI图像分类器:PyTorch实战Doompop项目全流程 ★ FEATURED ARTICLE
这年头搞技术的人很难绕开“AI”这个词。但说实话打开各种教程,迎面而来的不是十分钟带你掌握深度学习,就是三十天从入门到放弃的劝退长文。我见过太多人卡在第一步——他们装好了环境,跑通了官方Demo,但心里始终没有真正建立起“我自己也能搞一个AI项目”的底气。我自己也是从那个阶段过来的。所以当我准备开启一个新的小项目时,刻意给它起了一个很复古的名字:“Hello World, AI Doompop”。Hello World 是老程序员刻进DNA的仪式感,而 Doompop 是我给自己这个实验项目起的代号,听起来像是一只电子宠物或者一台末日风格的复古游戏机。这个名字其实宣告了一件事:这一次,我要把 AI 从“看别人玩”变成“自己亲手搭一个完整的东西”。这篇文章就是整个过程的完整记录,包括思路、踩坑、拆解和最终的参数、代码、数据是怎么落到实处的。如果你想找一个能真正练手的 AI 项目,跟着把这个流程走一遍,会比单纯刷一堆课程有用得多。1. 项目整体设计与思路拆解很多新手碰到的第一个陷阱,是选了一个宏大到自己根本驾驭不了的方向。什么“基于深度学习的XX系统”“多模态大模型XX平台”,听上去很酷,实际动手不到一周就发现算力不够、数据没有、代码看不懂,最后只剩一个空荡荡的 README 文件。我给自己定下的 Doompop 项目原则非常朴素:把范围缩到极小,把流程走完整。整个项目只做一件事——构建一个能识别“Doompop 角色是否处于危险状态”的图像分类器。听起来很简单对不对?但背后牵扯的东西其实一样都不少:数据采集、图像预处理、模型选型、训练调参、评估指标、可视化展示。把这些全部走一遍,你对 AI 项目全流程的理解,会比看十篇综述都深刻。1.1 核心需求解析这个项目要解决的核心问题非常具体:有一组以 Doompop 为主角的游戏渲染图,我们需要在每次渲染时自动判断画面中的 Doompop 是“安全”(周围无威胁)还是“危险”(被敌人包围)。这本质上是一个二分类图像识别任务。我为什么选择这个任务而不是热门的物体检测或者语义分割?理由有三个:二分类是图像识别的最小区块,模型结构、训练逻辑、评估方式都能讲透,不会因为任务本身的复杂度干扰你对主流程的理解。数据构造可以自行掌控,你不需要去网上爬取几万张实拍图,用基础图形和角色素材就能拼出可用的数据集,大大减少了入门的数据门槛。结果一眼可见,训练完之后,拿一张新渲染的图片丢进去,模型给出的是安全还是危险,你能直接感受到“AI确实学会了点什么”的成就感。1.2 为什么不用现成预训练模型直接套我知道有个很常见的思路:直接调一个训练好的 ResNet18,在 ImageNet 上做迁移学习,几行代码就能出结果。这也是现实中工业界最常见的做法。但我在 Doompop 这个项目里故意没有这么干,至少在第一版训练中没有直接用现成预训练权重。原因是我希望完整的训练逻辑在自己的掌控范围内。迁移学习固然快,但对于新手来说,它把太多东西封装成黑盒,你调了学习率、换了数据增强,却不知道底层到底发生了什么。而自己从头搭一个轻量的卷积网络,哪怕结构简单,你也能真真切切地感受到每一层张量形状的变化、每一个参数对训练结果的影响。说白了,第一遍做项目,不要怕从零开始,复杂模型以后有的是机会用。2. 环境准备与工具选型的权衡正式动手之前,环境的准备往往是第一个劝退点。很多人卡在依赖安装上,组合装不好,版本冲突一大堆。我在这个项目里踩过一轮坑之后,总结出了一套对新手最友好的组合。2.1 完整环境依赖清单为了保证自己复现的时候不出现莫名其妙的问题,我列了一个非常明确的依赖清单:操作系统:Windows 10/11 或 Ubuntu 20.04 均可,本项目对系统没有特殊要求Python:3.9 或 3.10,尽量避免最新的 3.13,部分依赖还不完全兼容PyTorch:2.1.0 或更新的 2.x 版本,CUDA 版本视显卡而定torchvision:与 PyTorch 版本配套,用于图像数据处理OpenCV-Python:4.x,用于图像预处理和可视化Matplotlib:3.x,用于绘制训练曲线和展示结果NumPy:1.24,矩阵运算基础库安装命令我一直用官方推荐的 pip 方式。这里有一个血的教训:千万不要直接从官网默认源安装 PyTorch,因为默认源在国内环境下那下载速度真的会让人崩溃。我是这样解决的:# 先安装 CPU 版 PyTorch 的依赖(不要急,下面是 GPU 版) # 如果是 N 卡且有 CUDA 环境,用以下命令安装 GPU 版 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 --mirror https://mirrors.cloud.tencent.com/pypi/simple那如果手里没有 N 卡怎么办?别慌,我用 CPU 也跑通过这个项目。只是训练时间会长一点,后面我会讲怎么基于 CPU 环境调整参数,照样能复现结果。2.2 为什么 PyTorch 而不是 TensorFlow作为一个入门项目,选深度学习框架也是个关键决策。我承认 TensorFlow 在工业界部署时依然有它的地位,但在“快速上手、逻辑清晰、调试方便”这三个维度上,PyTorch 的优势显而易见。PyTorch 的动态计算图是一个对新手极其友好的特性:你写代码的时候就像写普通的 Python 函数一样,中间的任何一行都可以打印张量的形状和值,非常方便排查问题。相比之下,TensorFlow 的静态图机制(尤其是旧版本)在你定义完计算图之后,再想去观察中间的某个值就非常别扭,需要 Session 那一套东西,学习曲线直接拉满。另外,PyTorch 在学术界的普及率已经是绝对的主流,这就意味着你遇到任何问题,几乎都能在社区里找到已有的讨论和解决方案。踩坑成本低,对新手来说比什么都重要。3. 数据集的构建:这部分往往决定成败很多教程讲 AI 项目,数据部分都是翻页带过的——“我们使用公开数据集”。但对一个真正要练手的项目来说,数据是自己亲手攒出来的,这个过程的价值完全不亚于模型训练本身。3.1 从零制作一个小而完整的图像数据集Doompop 项目的目标分类是“安全”和“危险”两种状态。我实际去构造数据的时候,总共生成了 1600 张图片,安全类和危险类各 800 张。为了不让模型偷懒只学到背景颜色之类的低级特征,我在数据生成阶段就故意加入了大量的随机化处理:角色位置随机:每张图片里 Doompop 不一定在正中心,可能偏左、偏上,甚至稍微小一点背景干扰随机:背景色在深色系和浅色系之间随机切换,再加一定的高斯噪声危险元素随机:如果属于“危险”类,就在角色周围随机添加红色警戒图标、敌方剪影或扩散的包围圈这里用到了最传统的 OpenCV 绘图就能实现,不需要任何高深的 CG 技术。每张图片统一下发到 128x128 像素,大小非常可控。生成完的图片按文件夹分好:data/ ├── train/ │ ├── safe/ │ │ ├── 0001.png │ │ ├── 0002.png │ │ └── ... │ └── danger/ │ ├── 0001.png │ ├── 0002.png │ └── ... └── val/ ├── safe/ │ └── ... └── danger/ └── ...训练集和验证集的比例是 8:2,即训练集 1280 张,验证集 320 张。这种简单的目录结构也是 PyTorch 的 ImageFolder 可以直接读取的格式,非常方便。3.2 数据增强到底要不要用我的经验是:对于这个小项目,数据增强不仅要用,而且要大胆地用。因为 1600 张图本身数量不大,模型很容易出现过拟合——也就是在训练集上表现好,一到验证集就掉链子。我在代码里加入了这些增强手段:from torchvision import transforms train_transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转 transforms.RandomRotation(degrees15), # 随机旋转15度以内 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # 使用标准归一化 ]) val_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意验证集是绝对不做随机增强的,否则你评估出来的指标根本没办法反映模型的真实水平。增强只针对训练集,这个原则不要搞混。3.3 数据集质量是真正的“胜负手”我见过太多人忽视数据质量,总想着“我模型结构够强,数据差点无所谓”。这个想法在图像分类任务里特别坑。如果只是简单地把角色放在纯色背景中间,模型可能只学到“背景颜色 危险程度”这种完全扭曲的规律,看起来训练指标很高,实际上毫无泛化能力。我第一次做这个项目时就犯了类似的错误。当时为了图省事,把“危险”类图片的背景都设置成了深红色。结果训练完模型一测,发现准确率虽然到了 95% 以上,但只要把一张背景偏红的“安全”图输入进去,模型就错判成危险。后来增加了背景随机化,重新生成了数据集,模型的泛化能力才正常起来。注意:图像分类模型非常容易在训练时“抄近道”,它会优先捕捉最容易区分的低级特征。如果这些特征跟真实的类别标签没有因果关系,那模型学到的就是伪规律。数据构建阶段一定不要偷懒,随机化越充分,模型学到的东西越真实。4. 模型结构搭建与训练流程实现数据准备完毕后,接下来进入整个项目最核心的实操环节:模型搭建和训练。这一步会让你的电脑真正“跑起来”,也是最容易出现各种问题的阶段。4.1 设计一个极简卷积网络关于模型结构,很多人的直觉是“越复杂越好”,但对于 128x128 这种小尺寸图片的二分类任务,一个太深太宽的网络只会让训练速度变慢,甚至出现过拟合。我最终的模型结构非常简单:import torch.nn as nn class DoompopClassifier(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), # 128x128x16 nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 64x64x16 nn.Conv2d(16, 32, kernel_size3, padding1), # 64x64x32 nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 32x32x32 nn.Conv2d(32, 64, kernel_size3, padding1), # 32x32x64 nn.ReLU(inplaceTrue), nn.MaxPool2d(2) # 16x16x64 ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 16 * 16, 128), nn.ReLU(inplaceTrue), nn.Dropout(p0.3), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这个结构里有几个关键设计值得解释。第一,连续的卷积层都加了 padding1,这样尺寸不会因为卷积而缩小得太快,信息保留更完整。第二,每个卷积之后都跟着一个 ReLU 激活函数,如果没有激活层,多层卷积叠加在一起就跟一层线性变换没有区别,模型的表达能力会非常弱。第三,最后一个分类器里加了一个 Dropout,比例是 0.3,它会在训练时随机丢弃 30% 的神经元连接,起到防止过拟合的作用。你可能会问:为什么没有用 BatchNorm?我的考虑是对于这种小而稳定的数据集,BatchNorm 带来的收益不明显,反而可能因为 batch size 太小引发一些不稳定的问题。小项目参数越少,每调一个变量就越能看清它的影响。4.2 训练参数的计算与选择逻辑搭建完模型之后,就算进入了训练环境的配置环节。有几个关键参数我在这里特别说明一下选定的逻辑,而不是直接甩给你一组数字。Batch Size。我最终设定了 32。为什么不设成 64 或者 128?原因有两点:一是受限于内存/显存,二是对于小数据集,较小的 batch size 反而往往带来更好的泛化效果,因为每次更新的梯度有更多随机性,相当于一种隐式的正则化。如果你的内存比较小,调成 16 也完全 OK。学习率。初始学习率设成 0.001。这个选择主要基于 Adam 优化器的特性——它对学习率的敏感度相对较低,0.001 几乎是它的“黄金默认值”。如果你用的是 SGD 优化器,那学习率通常要调大一些,比如 0.01 甚至 0.1,但相应地要配合动量(Momentum)才能稳定训练。训练轮数(Epochs)。我设了 30 轮。对于这么简单的二分类任务,30 轮已经足够看到模型收敛,再多轮数往往是在浪费算力。而且在训练过程中,我加入了学习率调度,每 10 轮就把学习率降为原来的十分之一,这样后期参数更新的步长更小,有助于让 loss 落到更优的极小值。4.3 完整训练脚本的落地实现下面是我实际跑通的训练脚本核心部分,每一步都加了注释:import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder from torchvision import transforms # 数据集路径 train_dir data/train val_dir data/val # 数据加载 train_dataset ImageFolder(train_dir, transformtrain_transform) val_dataset ImageFolder(val_dir, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers2) # 模型、损失函数、优化器 model DoompopClassifier(num_classes2) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 学习率调度器:每10轮衰减为原来的1/10 scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) # 训练循环 def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) avg_loss total_loss / total acc correct / total return avg_loss, acc # 验证循环 def evaluate(model, loader, criterion): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) avg_loss total_loss / total acc correct / total return avg_loss, acc # 主训练流程 num_epochs 30 best_val_acc 0.0 for epoch in range(1, num_epochs 1): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer) val_loss, val_acc evaluate(model, val_loader, criterion) scheduler.step() print(fEpoch {epoch:02d}/{num_epochs} fTrain Loss: {train_loss:.4f} Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f} Acc: {val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - Model saved at epoch {epoch}, val_acc {val_acc:.4f})里面有两个特别容易被忽略但又特别重要的细节。第一个是model.train()和model.eval()的切换。训练模式下 Dropout 会生效,但评估时必须切换到 eval 模式,否则验证时也会随机丢弃神经元,导致结果不稳定。这个点我在早期做项目时经常忘记,结果验证准确率飘忽不定,还以为模型写错了。第二个是torch.no_grad()上下文管理器。验证时不需要计算梯度,包一层这个可以显著减少内存占用、加速计算。别小看这个优化,数据量大了之后,这个差异非常明显。4.4 训练过程的“现场实录”我实际跑下来的结果大概是这样一个曲线:前 3 轮:训练 loss 从初始的 0.7 左右快速下降到 0.3 左右,验证准确率从 70% 左右蹿升到 90% 以上第 4 到第 10 轮:训练 loss 持续下降,但验证准确率在 92%~94% 之间波动第 11 轮(学习率衰减后):验证准确率明显上了一个台阶,稳定在 96% 以上第 20 轮(第二次衰减后):验证准确率最高到了 97.8%这个曲线是非常典型的“学习率衰减带来二次提升”的模式。如果你发现训练结束后验证准确率一直卡在同一个水平不涨,可以先尝试增大轮数并配合学习率衰减,往往比你调整模型结构更有效。5. 常见问题与排查技巧实录这个环节往往是博客里最值钱的部分。因为项目出 bug 时的排查经验,是任何教程都不会系统告诉你,但新手又一定会碰到的东西。我把自己在 Doompop 项目中实际碰到的典型问题列出来,每个都附带了排查思路和解决方式。5.1 验证准确率极低但训练准确率很高这种情况几乎可以断定是过拟合。我第一版代码里遇到过:训练准确率达到 99%,验证准确率却只有 78%。排查思路按优先级走:最优先检查数据会不会泄漏——比如训练集和验证集里出现了完全一样的图片。因为我是脚本批量生成的,有概率重复。这时直接检查两张集里的文件哈希对比即可确认。然后看数据增强是否设置正确,训练集增强和验证集增强是否出现混用。再看模型容量是不是过大——如果参数数量远超数据量,会把训练数据里的噪声和随机细节都背下来。Dropout 的比例是否太小或者根本没加。在做完数据去重、把增强策略分开、加上 Dropout 后,验证准确率直接拉到了 94% 以上。这个排查路径对大部分新手项目都适用。5.2 Loss 下降到一定程度就不再下降了loss 卡在某个水平下不去,未必是 bug。我在这个项目里就遇到 loss 在 0.1 左右卡了 5 个 epoch,验证准确率也纹丝不动。这时候我先确认了优化器和学习率设置没有明显问题,然后做了一个关键操作:把学习率衰减节点提前。在原始设置里,第一次衰减在第 10 轮,但实际到第 7 轮时 loss 已经趋于平缓。后来我改成每 7 轮衰减一次,卡顿现象消失,loss 继续下降到 0.04 左右。如果你不想频繁改调度器参数,另一个常见的做法是采用“ReduceLROnPlateau”调度器,它可以自动检测验证 loss 不下降的时机,然后手动降学习率:scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.1, patience3 )注意这种调度器需要在每个 epoch 结束时传入验证 loss 值,比如scheduler.step(val_loss),不要忘记这一步。5.3 CPU 环境下训练实在太慢怎么解决如果你只有 CPU,训练这个 30 轮的模型可能要跑 20 到 40 分钟,确实挺难受。但实际上手之后我发现了几个非常实用的提速思路:把模型输入图片尺寸从 128x128 降到 96x96,准确率几乎不受影响,速度提升 30% 以上减少通道数:把第一个卷积的 16 通道改成 12,甚至 8,参数量直线下降DataLoader 的 num_workers 调成 2(Windows 下不要超过 4),否则数据读取代价可能超过计算代价少量轮数验证模型是否正常收敛,比如先跑 3 个 epoch,确认 loss 确实在下降,再跑完整 30 轮,避免无效等待如果你的机器配置确实很低,也可以用 Google Colab 这种云端方案,把数据打包上传、跑完训练再下载模型,至少能免费借到一块 GPU。我的模型最终在 CPU 环境下也运行到了 96% 验证准确率,所以这个项目本身不依赖高端硬件。5.4 加载模型预测时总是报维度错误训练时好好的,但在加载模型做单张图片预测时报错,这个问题我碰到过两次,十有八九是输入图片的尺寸或通道数不对应。模型在训练时接收的输入张量形状是[batch32, channels3, height128, width128],而做单张预测时,如果你直接加载一张 128x128 的图片,得到的是[height, width, channels]的 NumPy 数组,没有 batch 维度。你需要这样处理:from PIL import Image import torchvision.transforms as transforms def predict_single_image(image_path, model, class_names): model.eval() img Image.open(image_path).convert(RGB) # 确保是三通道 transform transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img_tensor transform(img).unsqueeze(0) # 增加 batch 维度 with torch.no_grad(): outputs model(img_tensor) _, predicted torch.max(outputs, 1) return class_names[predicted.item()].unsqueeze(0)这一步就是给数据增加一维 batch 维度,这是新手最容易漏的。另外注意图片一定是 RGB 三通道,如果输入是灰度图,需要先转换成 RGB。5.5 训练结果无法复现代码没改,但每次跑出来的验证准确率都不一样,浮动在 1%~3% 之间。这个问题看起来不大,但确实影响项目的可复现性和调参体验。根源在于随机性。模型初始化的权重、数据的 shuffle 顺序都引入了随机因素。要想结果稳定,在训练代码主入口设置随机种子即可:import torch import numpy as np import random def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) # 这里设置 CUDNN 为确定性模式,CUDA 下运行时性能会稍微受影响 torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False set_seed(42)设置随机种子这件事虽然简单,但它是让项目从“能跑”变成“能复现”的关键一步。包括我后面调参时,也全靠这个种子固定了基线,不然每个参数改动的效果都会被随机噪声干扰,根本没法准确判断哪个变量真正起了作用。6. 从 Demo 到项目:还能怎么扩展Doompop 这个项目我可以明确说,它是一个“小但完整”的闭环:有明确的问题定义、有自己构造的数据集、有一个能跑通完整训练流程的模型、有评估指标、也有保存和加载推理的路径。这已经比很多只跑通教程代码的项目有含金量得多。但如果你不满足于此,我给几条循序渐进的扩展方向,算是给这个项目留出的成长空间。第一个方向是把二分类变成多分类。不要只区分“安全”和“危险”,可以把状态细分成“安全”、“被观察”、“被追击”、“包围”等等级更细的状态。你会发现,分类的粒度一变,数据集的构造难度、模型输出的设计、评估指标的解读全都跟着变了,这是一个很好的进阶练习。第二个方向是给模型加一个“解释性”的输出。比如用 Grad-CAM 类的方法可视化模型到底关注图像的哪些区域。你可能会发现模型判断某个角色“危险”主要是因为它看到了一团红色图形,这个信息反过来能让你改进数据质量。第三个方向是把模型部署到一个真实的端侧环境中。比如写一个基于摄像头实时画面的推理脚本,当画面中的 Doompop(打印的实体卡牌或玩偶)被识别为“危险”时,自动触发一个警报动作或语音提示。这一步会把“训练脚本”变成“应用产品”,对你理解 AI 工程化的完整流程帮助极大。我个人在实际操作中的体会是,像 Doompop 这种规格的小项目,最适合用来建立“AI 项目手感”。所谓手感,就是你拿到一个问题时,能快速判断数据从哪里来、模型用什么结构起步、训练完该看哪些指标、坏结果大概率坏在哪个环节。这些事情没做过三五个项目,很难从书里真正学到。一个完整的迷你项目,哪怕再小,也比十篇零散教程更能帮你建立这套判断力。最后再分享一个做这类项目的小技巧:每一个关键实验结果出来之后,记得随手记录下来,哪怕只是几句话。比如“加了背景随机化后,验证准确率涨了 12%”“学习率衰减从 10 轮改为 7 轮后,loss 突破 0.1”。等你积累足够多的记录,再回头审视这些数据,你对“哪些改动真正有效”的理解会变得非常深刻。这比任何模型调参指南都更贴近你手上项目的真实情况。
阅读完成 · 觉得有帮助?
咨询建站