上手复现经典网络的时候遇到的第一座山往往就是AlexNet。明明结构看起来不复杂真到了自己拿PyTorch从零写一遍卷积核大小、padding到底取多少、全连接层怎么接、训练时loss怎么死活降不下去问题一个接一个。这篇文章就把我在PyTorch里从空文件写到能跑通训练、又能稳定收敛的完整过程拆开来讲包括每一层的参数推导、代码怎么组织、训练推理和踩坑记录全程会给出可以直接复制运行的完整代码块。这篇内容适合三类人刚学完PyTorch基础语法、想拿一个完整项目练手的人课程作业或论文复现需要快速搞定AlexNet的人以及已经在跑现成代码、但遇到训练不收敛或显存爆掉等异常、想搞明白根源的人。我尽量按实际动手的顺序来讲而不是按文档目录的顺序这样读下来你会更清楚每个环节为什么要那么写。1. 动手前的整体思路别急着写代码先把结构算清楚1.1 AlexNet四个关键创新点以及复现时如何取舍AlexNet之所以是里程碑不是因为它网络层数深而是它把当时分散的几项技术组合到了一起并且用GPU把规模撑起来了。对我来说复现的价值在于理解“为什么每个设计点会出现在那个位置”。第一点是ReLU激活函数。在AlexNet之前主流是tanh和sigmoid但梯度在深层网络中非常容易饱和。ReLU在正区间的导数是常数1反向传播时梯度传递更稳定训练速度也快不少。这点在PyTorch里就是一个nn.ReLU(inplaceTrue)没有实现成本但你要理解它放在卷积和全连接之后的意义。第二点是Dropout。AlexNet的FC层有4096维参数量非常大训练时很容易过拟合。Dropout在训练时随机把一部分神经元置零迫使网络不依赖某个特定神经元相当于在一个Batch里做了多模型的隐式集成。PyTorch里是nn.Dropout(p0.5)要注意的是推理时必须关了它PyTorch的model.eval()会自动搞定但如果你自己写推理循环忘了切模式就麻烦了。第三点是LRN局部响应归一化。这个在现代复现里我基本不建议用PyTorch虽然有nn.LocalResponseNorm但大量实验证明它对最终准确率的提升非常有限而随着BatchNorm的出现LRN已经被完全取代。所以我会在后面代码里用BatchNorm替代LRN的位置效果更好收敛也更快。第四点是重叠池化。AlexNet用的MaxPool是3x3核、stride2相邻池化窗口有重叠。相比不重叠的2x2窗口重叠池化能轻微缓解过拟合这一点在CIFAR-10级别的数据上影响不大但代码里我会保留原始设计因为这也影响后面特征图的尺寸计算。1.2 选择PyTorch而不是其他框架的原因复现AlexNet这种事用PyTorch明显比TensorFlow顺手。倒不是说谁强谁弱而是PyTorch的nn.Module设计让“网络结构”和“前向传播逻辑”高度统一写模型就像在纸上画流程图。想打印中间特征图的形状直接在forward里加一行print(x.shape)就行想调试某个层的梯度torch.autograd的机制让人很容易介入。另外PyTorch的动态计算图特性在处理“按条件改变网络结构”这类需求时非常灵活虽然AlexNet这种纯前馈网络用不到但你后面做注意力机制、做检测头的时候会体会到这个优势。再加上社区里HuggingFace、Ultralytics等主流库都是PyTorch生态我用它积累的经验可以平移复用而不是学一套框架绑定一套思路。1.3 复现前的网络尺寸和显存预算在下手写代码前我习惯先把网络每一层的输出尺寸在草稿纸上推一遍。卷积层输出尺寸公式是输出尺寸 floor((输入尺寸 - 卷积核大小 2 * padding) / stride) 1原始AlexNet输入是224x224的ImageNet图像第一层卷积核11x11、stride4、padding2代入公式得到55池化后变成27后面每层以此类推。但如果你只是想跑通流程CIFAR-10的32x32输入直接套原结构是行不通的因为第一层卷积后尺寸变化太快后续的55x27这种尺寸对不上了。最常见的做法是给网络加一个nn.AdaptiveAvgPool2d((6, 6))让全连接层之前的特征图自适应到固定大小这样无论输入是32x32还是224x224都能接上后面的FC层。这个方案虽然不是最“纯正”的AlexNet但工程上最稳妥。另一个要提前算的是显存。用224x224输入、batch size为64训练原始AlexNet在单张8GB显存的卡上会非常勉强。Layer1的卷积输出是64x96x55x55这一张特征图就要649655554字节约74MB反向传播存梯度还要翻倍。所以我在实际实验里常用CIFAR-10以减轻显存压力后面如有余力再换大图。2. 数据准备与预处理好数据是训练稳定的前提2.1 数据集选型ImageNet太大CIFAR-10刚好能跑通全流程原始AlexNet是在ImageNet的120万张图片上训练的这个规模个人电脑几乎跑不动。我建议第一步先用CIFAR-10它只有10个类别6万张32x32的图下载下来也就160MB左右单张消费级显卡十几分钟就能跑一个像样的结果。CIFAR-10虽然分辨率低但用来验证模型代码是否正确、训练流程是否通顺完全够用。等你在小数据上跑通了再切换到大尺寸数据集只需要改数据集路径和输入尺寸模型本身改动很小。如果你的目标就是为了在ImageNet上出结果那确实需要多卡训练和数据并行这个后面我会简单提一下但作为完整代码解析我默认场景是单卡跑CIFAR-10。2.2 DataLoader和归一化参数的选择逻辑图像数据进网络之前归一化是一件绕不开的事。CIFAR-10数据集的像素值本来是0到255的整数我先把它们除以255变成0到1的浮点数再用每个通道的均值和标准差做标准化。CIFAR-10通用的均值和标准差是(0.4914, 0.4822, 0.4465)和(0.2470, 0.2435, 0.2616)这组数据是社区统计出来的直接用就行。选用torchvision.datasets.CIFAR10时transform参数里把ToTensor和Normalize串联起来。注意ToTensor已经自动做了除以255的操作所以不要在代码里再手动除一次。import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) train_dataset datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) test_dataset datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_size128, shuffleFalse, num_workers4, pin_memoryTrue)shuffleTrue只在训练集上打开测试集不需要因为测试时我们要逐个样本评估打乱顺序没有任何意义。num_workers可以根据CPU核数调整Windows上如果报错就减到0或者2。pin_memoryTrue在GPU训练时能加快数据从CPU内存拷贝到显存的速度这个在数据量大时收益更明显。2.3 数据增强策略RandomCrop和RandomHorizontalFlip为什么够用很多新手会纠结要不要上复杂的增强比如色彩抖动、随机擦除、Cutout。我的建议是在CIFAR-10上先做两个最基础的随机裁剪和水平翻转。RandomCrop(32, padding4)的做法是先把32x32的图pad到40x40再随机裁回32x32相当于给网络提供了一些平移不变性。RandomHorizontalFlip让图片有一定概率左右翻转对CIFAR-10里大多数物体类别都是合理变换可以让训练数据量翻倍。这两个增强的作用是缓解过拟合。AlexNet的参数量高达数千万直接在小数据集上硬训训练集loss很快会降到接近0测试集准确率却上不去这就是过拟合。增强相当于给网络提供了不同视角的样本让它学习到更鲁棒的特征。3. AlexNet网络结构逐层拆解与PyTorch代码实现3.1 网络整体参数一览表在写模型类之前我把每层配置整理成了一张表后面写代码时对照着填参数不容易犯错。层名称类型参数配置输出尺寸输入224x224参数量conv1Conv2d3-96, kernel11, stride4, padding255x55x96约3.5万pool1MaxPool2dkernel3, stride227x27x96无conv2Conv2d96-256, kernel5, padding227x27x256约61万pool2MaxPool2dkernel3, stride213x13x256无conv3Conv2d256-384, kernel3, padding113x13x384约88万conv4Conv2d384-384, kernel3, padding113x13x384约132万conv5Conv2d384-256, kernel3, padding113x13x256约88万pool3MaxPool2dkernel3, stride26x6x256无fc1Linear9216-40964096约3770万fc2Linear4096-40964096约1670万fc3Linear4096-10001000约409万如果把最后的1000类换成CIFAR-10的10类只需把fc3输出改成10。这个表里最值得注意的是fc1的输入维度9216它是pool3输出的6x6x256展平后得到的。如果用AdaptiveAvgPool这里算起来更省心。3.2 核心模型代码下面这段代码是完整可运行的AlexNet模型定义。我做了两个调整一是用BatchNorm替代LRN二是支持通过num_classes参数控制最后的分类数。import torch import torch.nn as nn class AlexNet(nn.Module): def __init__(self, num_classes10, in_channels3): super(AlexNet, self).__init__() self.features nn.Sequential( nn.Conv2d(in_channels, 96, kernel_size11, stride4, padding2), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size3, stride2), nn.Conv2d(96, 256, kernel_size5, stride1, padding2), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size3, stride2), nn.Conv2d(256, 384, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(384, 384, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(384, 256, kernel_size3, stride1, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size3, stride2), ) self.avgpool nn.AdaptiveAvgPool2d((6, 6)) 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), ) self._initialize_weights() def forward(self, x): x self.features(x) x self.avgpool(x) x torch.flatten(x, 1) x self.classifier(x) return x def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_in, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0)卷积层的输出尺寸用公式逐层验证一下224经过conv1变成(224 - 11 2*2)/4 1 55池化后变成27conv2保持27池化后13conv3、conv4、conv5都保持13最后池化成6。如果输入换成32x32的CIFAR-10conv1后尺寸是(32 - 11 4)/4 1 7池化后3后面几层维持3最后池化到1依然能通只是特征图太小了所以我推荐用AdaptiveAvgPool把它固定到6x6。3.3 权重初始化为什么重要这段代码里的_initialize_weights不是可有可无的装饰。PyTorch的nn.Conv2d默认初始化方式是Kaiming均匀分布但那种初始化是针对没有激活函数的情况设计的。AlexNet大量使用ReLUReLU会把负半轴的信息直接丢弃如果权重初始化不合理很多神经元一开始就输出为0梯度也一直是0这个神经元就永久死掉了。我采用Kaiming正太初始化它根据ReLU的特点计算了合适的方差范围让每一层输出的方差在传播过程中保持稳定。全连接层我用均值0、标准差0.01的正态分布这是AlexNet原文的做法。如果注释掉初始化函数跑一轮对比你会发现训练前期loss下降慢很多甚至可能出现loss变成nan的情况。4. 训练流程细节与超参数调优照抄不翻车4.1 损失函数与优化器配置分类任务的标准配置是交叉熵损失。PyTorch的nn.CrossEntropyLoss已经做了两件事先把模型最后一层的logits做softmax再计算交叉熵。所以模型最后一层不要额外加softmax否则会重复计算导致loss数值异常。优化器我选SGD加动量而不是Adam。很多新手会惯性用Adam但在图像分类这种需要精细调学习率的场景SGDmomentum的泛化能力通常更好。AlexNet原文用的momentum是0.9weight_decay设为5e-4。weight_decay就是L2正则它会让权重向量不会变得太大对抑制过拟合很有帮助。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)学习率0.01是相对保守的起点。如果你用的是batch size 256或更大可以尝试0.1但配套要多一些学习率衰减策略不然loss容易震荡。4.2 学习率调度StepLR还是CosineAnnealing训练迭代中学习率如果一直不变后期loss会在一个平台期来回波动很难继续下降。常见做法是训练到一定epoch后手动降低学习率比如每30个epoch缩小到原来的十分之一。PyTorch里用StepLR就能实现from torch.optim.lr_scheduler import StepLR scheduler StepLR(optimizer, step_size30, gamma0.1)step_size30表示每30个epoch降低一次gamma0.1表示乘以0.1。如果你的训练总轮数是90那就经历了三次衰减前30个epoch用0.0130到60用0.00160到90用0.0001。这个节奏比较符合AlexNet这种大模型的收敛习惯。另一种更省心的是CosineAnnealingLR它让学习率按余弦曲线从初始值平滑降到最低值不需要设定step_size天然适合Adam或SGD。我在实际对比中发现CIFAR-10上CosineAnnealing的最终准确率和StepLR几乎一样但前者的loss曲线更平滑少了一些冷启动阶段的波动。代码里两种都保留切换只需要改注释。4.3 训练循环一个结构清晰、方便扩展的框架训练循环写的质量直接决定你后面调试的效率。我把训练和验证拆成两个函数训练函数里每个epoch打印一次平均loss和准确率验证函数计算测试集上的Top-1准确率。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss 0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() avg_loss total_loss / total acc 100.0 * correct / total return avg_loss, acc def evaluate(model, test_loader, criterion, device): model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() avg_loss total_loss / total acc 100.0 * correct / total return avg_loss, acc num_epochs 90 best_acc 0.0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) test_loss, test_acc evaluate(model, test_loader, criterion, device) scheduler.step() print(fEpoch [{epoch 1}/{num_epochs}] fTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% f| Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%) if test_acc best_acc: best_acc test_acc torch.save(model.state_dict(), best_model.pth)这套模板里有一个很容易被忽略的关键点model.train()和model.eval()的切换。train模式下Dropout生效BatchNorm会更新running_mean和running_vareval模式下Dropout关闭BatchNorm用固定统计量。两个函数里都显式调用了对应模式防止验证时还把Dropout开着导致结果忽高忽低。4.4 训练时长与硬件选择建议在单张RTX 3060级别显卡上CIFAR-10训练90个epochbatch size为64大约需要25到35分钟。如果用纯CPU训练同样配置可能要三到六小时非常折磨。如果只有CPU我建议把batch size降到32同时把num_classes改为10减少全连接层输出的计算量并将epoch数缩短到30只验证流程是否正确不追求精度。想用GPU又没条件的人可以考虑云GPU平台按需租用或者用Kaggle、Colab提供的免费GPU额度。这些平台的环境基本都预装好了CUDA和PyTorch把代码贴进去就能跑。5. 推理与模型评估从训练到落地的最后一步5.1 加载模型做单张图片分类模型训练好之后最关键的是把它用于对单张图片的推理。做法是先加载state_dict再把模型切到eval模式最后写一个预处理函数把图片变成模型能接受的张量。from PIL import Image def preprocess_image(image_path): img Image.open(image_path).convert(RGB) transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) img_tensor transform(img).unsqueeze(0) return img_tensor CIFAR10_CLASSES [airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck] state_dict torch.load(best_model.pth, map_locationdevice) model.load_state_dict(state_dict) model.eval() image_tensor preprocess_image(test_cat.jpg) with torch.no_grad(): logits model(image_tensor.to(device)) prob torch.softmax(logits, dim1) pred_class torch.argmax(prob, dim1).item() print(fPredicted class: {CIFAR10_CLASSES[pred_class]}, probability: {prob[0][pred_class].item():.4f})map_locationdevice这个参数很重要。如果你在GPU上训练的模型要拿到CPU机器上推理不写这个参数直接load会报错报错信息通常类似Attempting to deserialize object on a CUDA device。写了map_locationcpu能解决这个问题。5.2 Top-1与Top-5准确率的统计方法ImageNet的比赛指标常用Top-5也就是说只要真实标签在模型预测概率最高的前5个类别里就算预测正确。CIFAR-10总共才10个类别Top-5的区分度太低通常是看Top-1。但如果后续你要在大类别数据集上做评估Top-1和Top-5的代码都需要写。def topk_accuracy(output, target, topk(1, 5)): maxk max(topk) batch_size target.size(0) _, pred output.topk(maxk, 1, True, True) pred pred.t() correct pred.eq(target.view(1, -1).expand_as(pred)) res [] for k in topk: correct_k correct[:k].reshape(-1).float().sum(0, keepdimTrue) res.append(correct_k.mul_(100.0 / batch_size)) return res输出的每个元素表示“前k个预测里包含真实标签的样本占比”。这套逻辑放在验证集上统计时要注意一次性不能塞太多Batch否则极少数样本的误差会被淹没我习惯按Batch汇总后取平均。5.3 模型保存的两种方式对比PyTorch保存模型有两种常见方式只保存state_dict或保存整个model。我的建议是始终保存state_dict。# 推荐只保存权重 torch.save(model.state_dict(), alexnet_cifar10.pth) # 加载时需要先定义好模型结构 model AlexNet(num_classes10) model.load_state_dict(torch.load(alexnet_cifar10.pth))保存整个模型虽然加载时不用重新定义网络但代码一旦升级旧模型文件很容易因为版本兼容问题加载失败。只保存权重的方式更稳定而且文件更小一个CIFAR-10量级的模型权重只有几十MB如果你用稀疏化或者量化的手段还能进一步压缩。6. 常见问题排查与踩坑记录一次性解决官方文档不教的事6.1 训练时loss出现nan的排查思路loss变成nan的情况我遇到不止一次原因大致分四种。第一种是学习率太大导致梯度更新幅度直接把权重推到了数值溢出区间解决办法是把学习率从0.01降到0.001试试。第二种是数据没有归一化原始像素值范围0到255如果直接输进网络经过多层卷积后激活值量级会变得很大为数值不稳定埋下隐患。第三种是权重初始值方差过大碰到Deep ReLU这种激活函数输出分布会畸形。第四种是标签问题比如标签里有负数或者超出类别数的值。排查手法我推荐先看梯度在loss.backward()之后遍历模型的参数打印grad的均值和标准差。如果发现某一层梯度是nan基本可以定位到那一层附近的问题。6.2 模型在测试集上准确率停滞不前的破局方法我最初用Adam训练时CIFAR-10准确率卡在78%左右上不去换SGD加动量后很快冲到82%以上。原因在于Adam的每个参数自适应学习率在训练后期可能过小没有足够动力跳出局部平坦区域。另外如果测试准确率比训练准确率低很多说明过拟合了优先给数据增强加料或者增大weight_decay。训练准确率本身也很低那通常是模型容量或优化问题。可以先把batch size减小看看单次更新是否更稳定也可以先拿一小部分数据比如1000张跑过拟合测试如果模型能记住这批数据说明代码流程没问题再回头处理优化策略。6.3 常见问题速查表现象可能原因解决方案Loss为nan学习率太大 / 数据未归一化降低lr至0.001检查NormalizeLoss不下降优化器选择不当 / 网络未加激活换SGDmomentum确认ReLU存在测试准确率远低于训练过拟合增强数据增强 / 增大weight_decayGPU显存不足batch size太大减小batch size或启用梯度累积验证时结果不稳定忘记model.eval()验证循环中显式调用model.eval()加载模型报Key不匹配num_classes与原模型不一致检查最后一层Linear的输出维度6.4 梯度累积让batch size可以“虚拟放大”如果你的显卡显存只能吃下batch size 32但实验设计需要batch size 128可以用梯度累积模拟大batch。做法是每4个小batch更新一次参数中间只累积梯度不执行optimizer.step()。accumulation_steps 4 optimizer.zero_grad() for i, (images, labels) in enumerate(train_loader): images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps这样才能保证累积后的梯度量级和真实大batch一致。这个技巧不改变模型代码只是在训练循环里做了节奏控制非常适合显存有限又不希望改动实验结果的人。7. 从CIFAR-10迁移到ImageNet的适配指南模型在CIFAR-10跑通之后换到更大数据集的改动比想象中少。首先是输入尺寸从32改成224这会影响网络第一层的卷积参数。原版AlexNet第一层用了11x11、stride4如果你直接从32x32的CIFAR-10切到224x224第一层代码不用改因为224输入正好对应原设计。但如果你的数据集图片尺寸是64x64或96x96建议把第一层的stride改小否则空间信息压缩太多。其次是类别数。ImageNet有1000类把num_classes改成1000即可输出维度变了但你保存的模型文件和CIFAR-10版本不兼容需要重新训练。最后是数据加载。ImageNet数据集文件夹比较大用torchvision.datasets.ImageFolder加载注意训练集和验证集要分开存放目录结构是train/类别名/图片.jpg和val/类别名/图片.jpg。再提一句数据并行。如果手头有多张显卡可以用torch.nn.DataParallel把模型包一层代码改动很小但BatchNorm在多卡场景下会有同步统计量的问题需要额外处理。更推荐的做法是用torch.distributed的DDP虽然是另外一个深坑但单机多卡的训练速度和稳定性都比DataParallel好。我个人在实际复现中的体会是写这个网络的过程其实是在复习卷积神经网络的底层逻辑。Conv2d的padding、stride和输出尺寸的关系你不亲手推一遍就永远只是背公式Dropout和BatchNorm在train和eval模式下的行为差异你不踩一次验证结果波动的坑就记不牢靠。AlexNet放到今天算力环境下已经算不上大模型但它承载的那些设计思想现在看依然是很多现代网络的起点。最后分享一个小技巧训练时每隔几个epoch用torch.save(model.state_dict(), checkpoint_epoch{}_acc{:.2f}.pth.format(epoch, test_acc))把中间结果存下来就算后边训练崩了至少有一个接近最优的备份能继续用这个习惯能帮你省下不少重跑时间。