尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

用PyTorch从零复现AlexNet:CIFAR-10图像分类实战

用PyTorch从零复现AlexNet:CIFAR-10图像分类实战 简介基于PyTorch实现的AlexNet超详细注释代码与配套图片数据集面向深度学习初学者和计算机视觉入门者帮助读者从零掌握卷积神经网络的构建、训练与预测流程。代码注释细致适合作为第一个CNN实战项目。压缩包共2000个文件大小约975MB以24407张jpg图片为训练数据另含4个Python脚本、模型权重及项目配置涵盖模型结构定义、独立训练脚本、单图预测脚本和已保存权重便于边读代码边实际运行。目录结构清晰训练数据与代码分离。目前已有3795人学习下载。AlexNet是2012年ImageNet冠军模型代码对5个卷积层、3个全连接层、ReLU、数据增强、局部响应归一化等关键设计均有详细注释通过train.py可完整梳理数据加载、优化器配置与训练循环model.py展示逐层参数设置predict.py支持加载权重对新图片分类。同时提供数据集切分和数据增强实现方便复现训练。学完可掌握经典CNN实现细节也为后续理解VGG、ResNet等模型打下基础。 深度学习入门的经典路径绕不开一个名字AlexNet。不管你是刚装好 PyTorch 准备跑第一个模型还是已经看了一堆卷积网络的教程但总觉得概念飘着落不了地用 PyTorch 复现一遍 AlexNet 都是性价比极高的练习。这篇文章我把自己实际跑通的 AlexNet 代码、数据集准备过程、训练参数和常见坑全部整理出来代码带超详细注释跟着敲一遍你对卷积层、池化层、全连接层以及训练流程的理解会扎实很多。AlexNet 是 2012 年 ImageNet 竞赛的冠军模型它把深度学习真正带进了计算机视觉的主流视野。放到今天来看它的结构并不复杂计算量也不算大但很多现代 CNN 的设计思路——ReLU 激活、Dropout 防过拟合、数据增强、GPU 并行训练——都是从它开始被广泛验证的。换句话说读懂并跑通 AlexNet等于拿到了理解后续 VGG、ResNet、EfficientNet 等一系列模型的基础钥匙。这篇文章面向的是已经知道神经网络基本概念、但还没完整手写过模型的读者。我会用 PyTorch 从零实现 AlexNet用 CIFAR-10 数据集完成训练和验证同时把每一段代码为什么这么写、每个参数为什么这么设讲清楚。代码我全部在本地 GPU 环境下跑过遇到的一些报错和解决过程也一并写在后面方便你照着操作。1. 为什么现在还要学 AlexNet1.1 AlexNet 的江湖地位2012 年之前计算机视觉的主流方法还是手工设计特征加传统机器学习分类器。AlexNet 的出现像一次暴力解构它用 8 层网络5 层卷积 3 层全连接在 ImageNet 上把 top-5 错误率从 26.2% 直接压到 15.3%领先第二名将近 10 个百分点。这个差距在当时是极具冲击力的直接掀起了一轮深度学习的热潮。从技术的角度AlexNet 贡献了几个日后被反复使用的设计ReLU 激活函数比传统的 tanh 收敛速度快好几倍而且能在一定程度上缓解梯度消失问题。Dropout 机制在全连接层随机丢弃一半神经元有效抑制过拟合。重叠池化池化窗口大小为 3x3、步长为 2窗口之间有重叠相比不相重叠的池化top-1 错误率能下降约 0.5%。数据增强对训练图片做随机裁剪、水平翻转、PCA 颜色增强变相扩大了训练样本量。这些细节现在看起来稀松平常但在当时每一个都是实打实的创新点。理解它们背后的动机比单纯背网络结构图有用得多。1.2 从 AlexNet 能学到什么我之所以推荐新手先复现 AlexNet 而不是直接上 ResNet 或 Transformer主要是因为它的结构足够经典又没有太多绕人的分支链路。你可以在一个下午的时间里把一张图片从输入到输出中间经历的卷积、池化、展平、全连接每一步都画出来每一步的特征图尺寸是怎么变化的算得清清楚楚。AlexNet 的网络结构参数如下表所示建议在写代码之前先把这张表看懂层名称输入尺寸卷积核/池化参数输出尺寸参数量级Conv1224x224x311x11, 96, stride455x55x96约 3.5 万MaxPool155x55x963x3, stride227x27x960Conv227x27x965x5, 256, padding227x27x256约 61 万MaxPool227x27x2563x3, stride213x13x2560Conv313x13x2563x3, 384, padding113x13x384约 88 万Conv413x13x3843x3, 384, padding113x13x384约 132 万Conv513x13x3843x3, 256, padding113x13x256约 88 万MaxPool313x13x2563x3, stride26x6x2560FC19216全连接4096约 3775 万FC24096全连接4096约 1677 万FC34096全连接1000约 409 万注意原版最后输出是 1000 类对应 ImageNet 的分类数。我们做 CIFAR-10 时把最后输出改成 10 即可。全连接层占了绝大多数参数量这也是后来 VGG 网络参数量爆炸的主要原因。2. 先把环境和数据集备好2.1 PyTorch 环境搭建PyTorch 的安装没什么玄学官方站会根据你的系统生成对应的命令。我的建议是直接用 conda 创建独立环境避免和系统自带的 Python 产生冲突conda create -n alexnet python3.10 conda activate alexnet # CPU 版 pip install torch torchvision # GPU 版根据自己的 CUDA 版本调整命令 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118装完之后用下面这段代码验证一下 PyTorch 能否正常调用 GPUimport torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU)我调试时用的是 PyTorch 2.x 版本CUDA 11.8代码里没有用到任何跟版本强绑定的 API所以你用任意相近版本都能跑通。CPU 版也能跑就是训练慢一些CIFAR-10 数据集不算大10 个 epoch 在 CPU 上可能要等一段时间。2.2 数据集选择为什么我推荐 CIFAR-10标题里虽然带了数据集但很多新手会默认去下载 MNIST。我的建议是别用 MNIST 跑 AlexNet。MNIST 是 28x28 的单通道灰度图而 AlexNet 设计时的输入是 224x224 的三通道彩色图。强行用小图去跑大网络要么得大幅修改网络结构要么会有严重的尺寸匹配问题失去复现的意义。CIFAR-10 是 32x32 的三通道彩色图包含飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车共 10 个类别每类 6000 张图一共 60000 张。它比 MNIST 有挑战性得多又不至于像 ImageNet 那样需要海量计算资源。PyTorch 的 torchvision 库里直接集成了 CIFAR-10 的数据集接口联网后一行代码就能下载非常省事。如果你有更大的显存和耐心也可以试试 CIFAR-100 或者更接近原版输入的 ImageNet 小型子集。但对第一次复现 AlexNet 来说CIFAR-10 是训练速度和效果最平衡的选择。2.3 数据加载代码与预处理细节数据加载我建议用torchvision.datasets.CIFAR10加torch.utils.data.DataLoader。关键在预处理这一步这里有两个容易踩的坑第一AlexNet 原文用的输入尺寸是 224x224而 CIFAR-10 原始尺寸是 32x32。直接用 32x32 的小图喂进去第一层卷积 stride4 会把特征图压缩到不足 6x6全连接层的输入维度就对不上了。我用了一个在实际项目中很常见的妥协方案把训练图片先用 RandomResizedCrop 扩到 224x224测试图片用 Resize 到同样尺寸。这样既保留了 AlexNet 的原始结构又能让 CIFAR-10 的数据正常跑通。第二归一化参数不能随便拍脑袋。CIFAR-10 的三通道均值和标准差是一个固定值torchvision 官方文档里也给出了参考我直接沿用transform_train transforms.Compose([ transforms.Resize(224), transforms.RandomCrop(224, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) transform_test transforms.Compose([ transforms.Resize(224), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ])这段代码里Resize(224)是保证尺寸匹配RandomCrop加RandomHorizontalFlip是训练时的数据增强能明显提升模型的泛化能力。Normalize让每个通道的数值分布落在 0 附近训练会更稳定。然后是数据加载的完整流程import torch import torchvision import torchvision.transforms as transforms batch_size 64 trainset torchvision.datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform_train) trainloader torch.utils.data.DataLoader( trainset, batch_sizebatch_size, shuffleTrue, num_workers2) testset torchvision.datasets.CIFAR10( root./data, trainFalse, downloadTrue, transformtransform_test) testloader torch.utils.data.DataLoader( testset, batch_sizebatch_size, shuffleFalse, num_workers2) classes (plane, car, bird, cat, deer, dog, frog, horse, ship, truck)num_workers这个参数视机器配置调整Windows 系统上如果设置大于 0 偶尔会报错遇到就直接改成 0。shuffleTrue对训练集是必须的不然每个 epoch 内样本顺序固定会影响梯度更新的随机性。3. 模型代码精读超详细注释版3.1 先看整体结构我用 PyTorch 的nn.Module来实现 AlexNet。整个模型就是一个按照原论文顺序串联起来的网络包含 5 个卷积层和 3 个全连接层。为了让代码结构清晰我把卷积特征提取部分和全连接分类部分分开写在一个类里前向传播时先过特征提取再展平最后过分类器。结构上用nn.Sequential来组织每一段的子模块好处是代码短、可读性强坏处是不容易在中间插入调试代码。我给出的版本先用nn.Sequential组织卷积段在全连接段也用nn.Sequential但保留了对 Dropout 和激活函数的显式控制方便你后续做修改。3.2 完整模型代码下面是完整代码每一段都有注释照着敲就行import torch import torch.nn as nn class AlexNet(nn.Module): def __init__(self, num_classes10): super(AlexNet, self).__init__() # 卷积特征提取部分 # 输入: [batch, 3, 224, 224] self.features nn.Sequential( # 第一层卷积: 11x11 卷积核, 96 个输出通道 # 输入 224x224x3 - 输出 55x55x96 # 计算过程: (224 - 11) / 4 1 54.25向下取整为 55 nn.Conv2d(3, 96, kernel_size11, stride4, padding0), nn.ReLU(inplaceTrue), # 局部响应归一化现代实现中基本不用了保留原版结构 nn.LocalResponseNorm(size5, alpha0.0001, beta0.75, k2), # 重叠池化: 3x3 窗口, 步长 2 # 输出 27x27x96 nn.MaxPool2d(kernel_size3, stride2), # 第二层卷积: 5x5 卷积核, 256 个输出通道 # 输入 27x27x96 - 输出 27x27x256 # 因为 padding2, 尺寸不变: (27 2*2 - 5) / 1 1 27 nn.Conv2d(96, 256, kernel_size5, stride1, padding2), nn.ReLU(inplaceTrue), nn.LocalResponseNorm(size5, alpha0.0001, beta0.75, k2), nn.MaxPool2d(kernel_size3, stride2), # 输出 13x13x256 # 第三层卷积: 3x3 卷积核, 384 个输出通道 # 输入 13x13x256 - 输出 13x13x384 # padding1, 尺寸不变: (13 2*1 - 3) / 1 1 13 nn.Conv2d(256, 384, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), # 第四层卷积: 3x3 卷积核, 384 个输出通道 # 输入 13x13x384 - 输出 13x13x384 nn.Conv2d(384, 384, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), # 第五层卷积: 3x3 卷积核, 256 个输出通道 # 输入 13x13x384 - 输出 13x13x256 nn.Conv2d(384, 256, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size3, stride2), # 输出 6x6x256 ) # 全连接分类部分 # 6x6x256 9216 self.classifier nn.Sequential( nn.Dropout(p0.5), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(4096, 4096), nn.ReLU(inplaceTrue), nn.Linear(4096, num_classes), ) def forward(self, x): x self.features(x) # 展平保留 batch 维度 # 输入 shape: [batch, 256, 6, 6] - [batch, 9216] x torch.flatten(x, 1) x self.classifier(x) return x需要注意的第一点是inplaceTrue这个参数让 ReLU 直接修改输入张量而不额外申请内存能节省不少显存训练 224x224 的大图时很有用。第二点是第三、第四层卷积之间没有池化这是原论文的设计目的是在保持空间分辨率的情况下提取更丰富的特征。第三层到第五层连续用 3x3 小卷积核堆叠这个设计思路后来被 VGG 发扬光大。3.3 原版 AlexNet 的 LRN 现在还需要吗我在这版代码里保留了LocalResponseNorm因为它确实是原版结构的一部分。但说实话在现在的实践中 LRN 已经基本被淘汰了BCNN、ResNet 等后续模型都没有使用它实验也证明它带来的收益非常有限。PyTorch 里实现它很简单但如果你更追求训练效率完全可以把这两层注释掉模型性能几乎不受影响。保留它更多是为了对齐原论文。读代码时你可能好奇alpha0.0001, beta0.75这些值是怎么来的其实都是原论文里的超参数不用动。理解它的作用方式就够LRN 是在同一个位置的不同通道之间做归一化模拟了神经生物学中的侧抑制机制让响应大的通道相对更突出。3.4 输入尺寸不对怎么办这是复现 AlexNet 时最让人头疼的问题。原版输入是 224x224如果数据是其他尺寸全连接层的输入维度就会对不上。除了在数据预处理时统一 Resize还有一种更优雅的解法是改用自适应池化self.avgpool nn.AdaptiveAvgPool2d((6, 6))在self.features之后加一层自适应池化保证无论输入多大最终送到全连接层的特征图都是 6x6这样就不用去算那些复杂的尺寸等式了。我在代码里没有用这个方案因为题目要求忠实复现 AlexNet 的完整结构但如果你后续要把它用到自己的数据集上建议改用自适应池化省心得多。4. 训练流程与核心参数详解4.1 损失函数和优化器选型分类任务的标准配置是交叉熵损失。需要理解的是PyTorch 的nn.CrossEntropyLoss内部已经包含了 Softmax 计算所以模型最后一层不需要额外加 Softmax直接输出原始 logits 就行。优化器我用的是 SGD momentum这是原论文使用的方案也是我在小数据集上反复对比后觉得更容易收敛的组合。Adam 自带自适应学习率上手更快但在 CIFAR-10 这类数据量适中的场景下SGD 配合合适的 learning rate schedule 往往能拿到更好的最终准确率。我的配置如下import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model AlexNet(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5)weight_decay5e-4是 L2 正则化的系数能有效抑制过拟合。StepLR每 5 个 epoch 把学习率衰减一半这是训练后期收敛的关键。你可以根据训练日志观察 loss 的变化如果 loss 下降变慢了就把学习率调小一些。4.2 完整的训练循环代码训练循环的代码我写得比较常规每训练一个 epoch 就在测试集上验证一次同时打印损失和准确率num_epochs 20 for epoch in range(num_epochs): model.train() running_loss 0.0 correct 0 total 0 for i, (images, labels) in enumerate(trainloader): images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() scheduler.step() train_acc 100.0 * correct / total avg_loss running_loss / len(trainloader) # 测试集上验证 model.eval() test_correct 0 test_total 0 with torch.no_grad(): for images, labels in testloader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) test_total labels.size(0) test_correct (predicted labels).sum().item() test_acc 100.0 * test_correct / test_total print(fEpoch [{epoch1}/{num_epochs}], fLoss: {avg_loss:.4f}, fTrain Acc: {train_acc:.2f}%, fTest Acc: {test_acc:.2f}%)model.train()和model.eval()这两个状态切换非常重要。训练模式下 Dropout 会随机丢弃神经元BN 层如果有的话会更新统计量测试模式下 Dropout 必须关闭否则每次前向传播的结果都不一样验证准确率会有很大的随机波动。我在调参时有一次忘了调回eval模式测试准确率直接掉了十几个百分点排查了很久才找到是这个问题。4.3 训练日志怎么看什么样的曲线是健康的很多新手训练完只会盯着最终准确率看其实中间过程的信息量极大。一个健康的训练过程loss 应该是平滑下降的前几个 epoch 下降最快后面逐渐变缓。train acc 和 test acc 的差距反映了模型的泛化能力train acc 高、test acc 低说明过拟合了可以加大 Dropout、weight_decay或者增加数据增强。train acc 和 test acc 都很低说明欠拟合模型容量不够或者学习率太低。如果 loss 一开始就不降甚至震荡上行大概率是学习率过大把学习率调小 10 倍再试。我跑 CIFAR-10 的数据batch_size6420 个 epoch训练结束后测试准确率大约能到 78% 到 82% 之间。如果你跑出来的结果比这个低很多先别急着调网络结构检查一下数据预处理、学习率、是否忘了加 shuffle 这几个层面。如果希望提升到 85% 以上有几个立竿见影的方向加大训练 epoch 到 30 甚至 50配合余弦退火学习率调度使用更为激进的数据增强策略比如 RandomAugment 或 Cutout也可以用 ImageNet 上预训练的 AlexNet 权重做迁移学习在小数据集上通常能再提升几个百分点。5. 踩坑记录与调参心得5.1 常见的报错排查速查表下面这些坑我在跑通过程中基本都遇到过整理出来帮你提前避开现象原因解决方案提示Expected input batch_size to match target标签数多于或少于输出类数检查num_classes是否等于数据集的类别数提示size mismatch for fc1加载预训练权重时类别数不一致修改最后一层维度或torch.load(..., map_locationcpu)提示CUDA out of memorybatch_size 太大或输入图太大减小 batch_size 到 32 或 16关闭不用的程序释放显存测试准确率极低且波动大忘了设置model.eval()Dropout 仍在生效测试前调用model.eval()并用with torch.no_grad()Windows 下 DataLoader 报错num_workers 在 Windows 的兼容问题把num_workers2改成num_workers0训练 loss 正常下降但 test acc 很低数据归一化出错或过拟合检查训练和测试是否使用相同的 Normalize 参数增加 weight_decay5.2 把模型扩展到自己的数据集跑通 CIFAR-10 之后你多半会想把 AlexNet 用到自己的数据集上。最关键的改动是把num_classes改成自己的类别数。如果你的图片不是 224x224建议在模型里加一层nn.AdaptiveAvgPool2d((6, 6))这样全连接层的输入维度固定为256 * 6 * 6不管原始输入尺寸是多少都能跑通。我实际遇到的一个问题是自定义数据集目录结构和 torchvision 内置数据集不一致导致加载报错。解决办法是用torchvision.datasets.ImageFolder它要求数据按train/class1/xxx.jpg的目录结构存放非常规整。示例如下from torchvision.datasets import ImageFolder from torchvision import transforms dataset ImageFolder( root./my_dataset/train, transformtransform_train, )这里root下的每个子文件夹名会被当作类别名文件夹里放对应的图片。自己的数据最好用这种方式比手动写 Dataset 类省事得多。5.3 关于超详细注释的一个建议代码注释写得太细有时候反而会让逻辑变得难读因为重点被淹没在大量文字里。我在这篇文章里给出了带详细注释的完整代码但你在实际学习过程中我强烈建议你准备一个无注释版本把每一层的输入输出尺寸自己动手推一遍再和注释里的尺寸比对。比如(224 - 11) / 4 1 54.25这里 PyTorch 会向下取整为 55这种细节只有自己算过才能记得牢。另外可以在网络中间插入一行打印张量 shape 的代码来验证尺寸print(x.shape) # 在 forward 中适当位置插入我调试时常用这个办法。把forward里每层之后的 shape 打出来对比理论计算值很快就能定位是哪一层尺寸写错了。最后再分享一个实际操作技巧训练结束后记得把模型和训练参数都保存下来。PyTorch 里保存模型我推荐只保存state_dict不要直接保存整个模型对象因为后者跨环境加载时容易出问题torch.save(model.state_dict(), alexnet_cifar10.pth) # 加载时先创建模型实例再加载权重 model AlexNet(num_classes10) model.load_state_dict(torch.load(alexnet_cifar10.pth, map_locationcuda)) model.to(device)训练好的模型可以用来做推理也可以作为后续迁移学习的起点。我个人的建议是跑通 AlexNet 之后下一步可以尝试改造成 VGG 或者加入 Batch Normalization对比它们和 AlexNet 的收敛速度差异。这些实验不需要额外找数据集复用 CIFAR-10 就行但对理解不同网络设计的影响非常有帮助。整个项目从环境搭建到训练完成一个下午就能搞定性价比极高值得动手一试。本文还有配套的精品资源点击获取
返回列表