1. 为什么还要手写AlexNet不只是考古1.1 一个绕不开的baseline前几天有位读者问我现在都是ResNet、ViT的时代了还有必要花时间看PyTorch实现AlexNet的完整代码吗我的回答是非常有必要。AlexNet是深度学习在计算机视觉领域引爆的起点它几乎是所有现代卷积网络的“母版”卷积层堆叠提取特征、池化降维、Dropout防止过拟合、ReLU加速收敛……你在任何一本深度学习教材里都会遇到它但真正用PyTorch把它从零写一遍、跑通训练和预测会发现很多论文里没写清楚的细节。这篇博文就把我自己的完整代码和踩坑过程整理出来适合三类人刚入门CNN想找第一个完整项目的人、准备拿经典模型做课程设计或论文baseline的人、以及想理解迁移学习到底在迁移什么的人。1.2 完整代码要解决的核心问题很多人以为手写AlexNet就是堆几个nn.Conv2d。其实完整代码至少包括六个部分网络结构定义、数据准备与预处理、训练循环、验证评估、模型保存与迁移学习。如果少了任何一环你跑起来就会遇到“数据维度对不上”“训练loss震荡”“验证集准确率忽高忽低”之类的问题。所以这篇不是只贴一个类定义而是从数据到模型到训练全链路都拆开讲每个关键点我都会解释为什么这么做以及哪些地方可以根据自己的数据集调整。2. 网络结构拆解从原始论文化到PyTorch代码2.1 AlexNet的原始设计并不复杂的五个卷积层AlexNet论文发表于2012年当时因为单块GPU显存不够作者把网络分成两条流水线放到两块GPU上并行。今天我们在PyTorch里实现时完全不需要复刻这个双流结构直接用标准卷积就可以。核心结构是5个卷积层加3个全连接层输入是224x224的RGB图像。第一层用11x11的大卷积核配合步长4快速降分辨率后面逐渐换成5x5和3x3的小卷积核通道数从96一路增加到384再到256。我把原始参数整理成了表格方便对照代码时一眼看清层输出尺寸卷积核/步长/填充后续操作Conv196 x 55 x 5511x11, stride 4, pad 0ReLU LRN MaxPool 3x3/2Conv2256 x 27 x 275x5, stride 1, pad 2ReLU LRN MaxPool 3x3/2Conv3384 x 13 x 133x3, stride 1, pad 1ReLUConv4384 x 13 x 133x3, stride 1, pad 1ReLUConv5256 x 13 x 133x3, stride 1, pad 1ReLU MaxPool 3x3/2然后是展平接三个全连接层4096、4096、1000类别数。这里有个细节第一个全连接层的输入维度是Conv5输出特征图展平后的256x6x6不是随便写的。2.2 在PyTorch中定义features部分直接上代码。我用的是PyTorch 2.x但这段代码在1.8以上的版本都能跑。先定义卷积特征提取部分import torch import torch.nn as nn class AlexNet(nn.Module): def __init__(self, num_classes1000, dropout0.5): super(AlexNet, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 96, kernel_size11, stride4, padding0), nn.ReLU(inplaceTrue), nn.LocalResponseNorm(size5, alpha1e-4, beta0.75, k2.0), nn.MaxPool2d(kernel_size3, stride2), nn.Conv2d(96, 256, kernel_size5, stride1, padding2), nn.ReLU(inplaceTrue), nn.LocalResponseNorm(size5, alpha1e-4, beta0.75, k2.0), 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), )这里需要注意两点。第一最后一个卷积层输入是384输出是256不是256到256。第二PyTorch中nn.LocalResponseNorm的默认公式是带窗口归一化的如果你要严格对照论文size取5、alpha取1e-4、beta取0.75、k取2.0基本就是原始设置。第三个卷积层和第四个卷积层之间没有池化只有到最后第五层结束才做一次MaxPool。2.3 定义classifier与初始化继续写分类器部分self.classifier nn.Sequential( nn.Dropout(pdropout), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplaceTrue), nn.Dropout(pdropout), nn.Linear(4096, 4096), nn.ReLU(inplaceTrue), nn.Linear(4096, num_classes), ) def forward(self, x): x self.features(x) x torch.flatten(x, 1) x self.classifier(x) return x关于初始化原始论文使用了标准差为0.01的高斯分布初始化卷积层和全连接层并用常数1初始化第2、4、5卷积层和全连接层的偏置其余偏置为0。PyTorch默认初始化在ImageNet这种大数据集上问题不大但如果你要在小数据集上训练最好还是做一下初始化。经验是下面这套def initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, 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)kaiming初始化是ReLU系网络的常用选择如果完全照搬论文的0.01高斯分布也行只是小数据集上kaiming给了我更稳定的表现。2.4 LRN的争议与替代方案花点篇幅单独说LRNLocalResponseNorm。它的思想是在一个通道附近区域内做归一化让响应大的神经元的输出相对更大提高模型的泛化能力。理论上很优雅但在实际工程里它带来的提升非常有限而且计算偏慢。所以后来的VGG、ResNet基本都放弃了它改用BatchNorm或者什么都不用。如果你是在自己数据集上做实验我建议保留LRN复现论文效果但也要知道可以把它替换成nn.BatchNorm2d。替换的时候注意LRN放在ReLU之后、MaxPool之前而BatchNorm通常放在卷积之后、ReLU之前位置不一样。完整代码里我保留LRN主要是为了忠实于原始结构。3. 数据准备与预处理让模型吃上干净的数据3.1 数据集目录与ImageFolderPyTorch里最省事的图片分类读取方式是torchvision.datasets.ImageFolder。它要求你的数据按这种目录组织data/ train/ class_a/ 001.jpg 002.jpg class_b/ 001.jpg ... val/ class_a/ ...每个子文件夹的名字就是类别名。ImageFolder会自动按文件夹名字母顺序生成从0开始的标签索引你不用手动维护标签字典。但要注意如果类别名是中文或者命名不规范建议先统一整理。我自己习惯在代码里额外打印一份class_to_idx防止后续做推理时把预测索引对应错类别。3.2 数据增强与归一化的搭配训练集和验证集的预处理必须分开写。这是新手最容易踩的坑有人直接拿验证集的数据增强方式去训练导致验证时图像被随机裁剪准确率怎么都上不去。训练集我用from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), 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.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里mean和std用的是ImageNet统计值。如果你的数据集不是ImageNet在迁移学习场景下用这个归一化依然没问题但从零训练时最好自己统计一下数据集的均值和标准差否则输入分布跟模型预期不一致训练很容易不稳定。Resize(256) CenterCrop(224)是验证集的标准组合保证所有图都是224x224且内容居中。千万不要在验证集用RandomResizedCrop那会让同一张图每次验证结果都不一样。3.3 加载数据与DataLoader参数下面是完整的数据加载代码我会把每个参数的理由也标出来from torch.utils.data import DataLoader batch_size 64 train_dataset datasets.ImageFolder(rootdata/train, transformtrain_transform) val_dataset datasets.ImageFolder(rootdata/val, transformval_transform) train_loader DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( val_dataset, batch_sizebatch_size, shuffleFalse, num_workers4, pin_memoryTrue, drop_lastFalse )解释几个参数shuffleTrue只在训练集用验证集不需要打乱否则每个epoch的评估顺序都在变影响调试。pin_memoryTrue在GPU训练时把数据锁页减少CPU到GPU的拷贝时间。num_workers根据机器CPU核数调整Windows上经常因为多进程序列化问题报错可以先设0运行正常后再调大。drop_lastTrue是防止最后一个batch太小导致统计异常但AlexNet没有BatchNorm这个参数影响不大留着只是习惯。如果你的数据集没有预先划分train/val可以用random_splittotal len(train_dataset) val_size int(total * 0.2) train_size total - val_size train_subset, val_subset torch.utils.data.random_split( train_dataset, [train_size, val_size] )但注意random_split不按类别比例分配类别不均衡时最好先从每个类里各取一部分组合验证集不要偷懒。4. 训练循环全解析每一步都在做什么4.1 超参数为什么是这些值训练前先确认几个关键超参数并理解它们的来源参数值说明batch_size64看显存8G显卡跑224x224可以到6412G可以试128base_lr0.01原论文用0.01配合SGD小数据集可降到0.001momentum0.9标准动量帮助跳出局部极小weight_decay5e-4L2正则抑制过拟合epochs50~90论文90小数据集一般50够用lr_step30每30个epoch衰减一次lr_gamma0.1衰减到原来的10%用SGDmomentum而不是Adam是AlexNet这类传统CNN在大数据集上的常见选择。SGD配合学习率衰减的泛化能力通常更好但缺点是对学习率敏感。如果你刚开始调试我建议先用Adam(lr1e-3)快速看模型能不能收敛没问题再切回SGD精调。很多人的第一个项目死在SGD学习率0.01上loss直接飘到NaN不是模型写错了是学习率太大。4.2 优化器与学习率调度完整代码import torch.optim as optim model AlexNet(num_classeslen(train_dataset.classes)) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD( model.parameters(), lr0.01, momentum0.9, weight_decay5e-4 ) scheduler optim.lr_scheduler.StepLR( optimizer, step_size30, gamma0.1 )关于类别数len(train_dataset.classes)会直接读ImageFolder的文件夹数量这样改数据集时不用同步改代码不容易出错。4.3 单epoch训练与验证循环训练循环是整篇代码的核心我把它拆成两个函数职责清晰也方便复用。训练部分def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) 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) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc这里有个细节loss.item()要在backward之后调用但必须在step之前还是之后无所谓只要在计算图释放前取值就行。统计时用loss.item() * images.size(0)再除以总样本数避免不同batch大小不均导致平均失真。验证部分torch.no_grad() def validate(model, loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc验证函数上加了torch.no_grad()这非常关键。如果不加验证过程会构建计算图显存消耗翻倍而且完全没必要。model.train()和model.eval()的切换只影响Dropout这类层AlexNet里有Dropout训练时随机丢弃验证时必须关闭所以eval()不能忘。整个训练主循环best_acc 0.0 start_epoch 0 for epoch in range(start_epoch, 50): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc validate( model, val_loader, criterion, device ) scheduler.step() current_lr optimizer.param_groups[0][lr] print(fEpoch {epoch1:03d} | ftrain_loss{train_loss:.4f} train_acc{train_acc:.4f} | fval_loss{val_loss:.4f} val_acc{val_acc:.4f} | flr{current_lr:.5f}) if val_acc best_acc: best_acc val_acc torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), best_acc: best_acc, }, best_model.pth)我习惯在每个epoch后打印一行包含训练loss、训练acc、验证loss、验证acc和学习率。训练时观察loss和acc的变化比观察任何TensorBoard曲线都直接。如果train_loss下降而val_loss不降甚至上升就是过拟合信号需要Dropout、数据增强或提前停止。4.4 模型保存与断点续训上面保存的是checkpoint而不是单纯state_dict目的是支持断点续训。如果你训练到第40个epoch断电了不用从头来checkpoint torch.load(best_model.pth, map_locationdevice) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) scheduler.load_state_dict(checkpoint[scheduler_state_dict]) start_epoch checkpoint[epoch] 1 best_acc checkpoint[best_acc]注意load_state_dict之前要确保模型结构和保存时一致尤其是num_classes不能变。map_locationdevice可以让你在GPU保存的模型加载到CPU时不出错反过来也一样。5. 评估与可视化不只是看准确率5.1 混淆矩阵找错都错在哪一类分类模型的验证集准确率只是一个数字真正要深入分析要看混淆矩阵。比如100个类别里整体acc有90%但你不知道是每个类都90%还是98个类满分、2个类全错。用sklearn一行就能算from sklearn.metrics import confusion_matrix import numpy as np all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds)然后你可以用matplotlib画热力图或者直接打印每个类别的precision、recallfrom sklearn.metrics import classification_report print(classification_report( all_labels, all_preds, target_namesval_dataset.classes, digits3 ))这个报告对类别不均衡的项目尤其有用。我见过一个项目准确率看着还行但细分后发现模型完全没学会某个小类别原因是那个类别在训练集里只有几十张图被大类别淹没了。这时候要么加数据要么给Loss加类别权重不是改模型结构能解决的。5.2 特征图可视化看看网络到底学到了什么分析CNN的一种有效方式是可视化第一层卷积核和中间特征图。AlexNet第一层有96个11x11的卷积核把它画出来你就能直观看到低频、高频、颜色边缘等滤波器。提取特征图可以用一个简单方法把features部分的输出拿下来。def visualize_feature_maps(model, image_tensor, layer_index0): model.eval() activations None def hook_fn(module, input, output): nonlocal activations activations output handle model.features[layer_index].register_forward_hook(hook_fn) with torch.no_grad(): _ model(image_tensor.unsqueeze(0).to(device)) handle.remove() feature_map activations[0] # [C, H, W] feature_map feature_map.cpu() # 取前16个通道拼图展示 from torchvision.utils import make_grid grid make_grid(feature_map[:16].unsqueeze(1), nrow4, normalizeTrue) return grid这里的核心是register_forward_hook它可以在forward过程中拦截某一层的输出而不修改模型。可视化用得最多的是第一个卷积层和最后一个卷积层前者看出来的是颜色和边缘后者看出来的是更抽象的语义特征。5.3 简单日志脚本不用打开TensorBoard也能监控TensorBoard当然好用但有时候你只是想在终端看到稳定的验证曲线。我通常直接把每个epoch的train_acc、val_acc写到一个txt文件训练结束后用几行代码画图with open(history.txt, w) as f: f.write(epoch train_acc val_acc\n) # 在每个epoch循环里加一行 f.write(...) # 画图 import matplotlib.pyplot as plt epochs [] train_accs [] val_accs [] with open(history.txt) as f: next(f) for line in f: parts line.split() epochs.append(int(parts[0])) train_accs.append(float(parts[1])) val_accs.append(float(parts[2])) plt.plot(epochs, train_accs, labeltrain_acc) plt.plot(epochs, val_accs, labelval_acc) plt.xlabel(epoch) plt.ylabel(accuracy) plt.legend() plt.savefig(history.png)这种方法比TensorBoard轻量得多也更贴近“跑实验看结果”的日常场景。如果你确实想看loss曲线把train_loss和val_loss也一起记录就行。6. 迁移学习站在预训练权重肩膀上微调6.1 加载官方预训练权重很多时候你不需要从零训练AlexNet。torchvision直接提供在ImageNet上预训练好的权重import torchvision.models as models model models.alexnet(weightsmodels.AlexNet_Weights.IMAGENET1K_V1) num_classes len(train_dataset.classes) model.classifier[6] nn.Linear(4096, num_classes)torchvision的AlexNet结构和我们上面手写的基本一致所以替换最后一个全连接层就行。注意索引是6因为classifier是一个Sequential里面0是Dropout1是Linear2是ReLU3是Dropout4是Linear5是ReLU6是最后的Linear。如果你不确定打印model.classifier看一眼最稳妥。在PyTorch 1.x时代这行代码是models.alexnet(pretrainedTrue)2.0以后换成了weights参数。pretrainedTrue现在会报警告旧代码要迁移一下。6.2 冻结与解冻两种微调策略迁移学习的第一步通常是把预训练特征提取器冻结只训练新加的分类头。原因是小数据集上如果全量更新所有参数很容易把预训练学到的通用特征学坏导致过拟合。冻结方法很简单for param in model.features.parameters(): param.requires_grad False # 只训练分类器 optimizer optim.SGD( model.classifier.parameters(), lr0.001, momentum0.9, weight_decay5e-4 )先用冻结方式跑几个epoch看验证集acc是否有上升趋势。如果上升明显再决定要不要解冻全部层做微调for param in model.parameters(): param.requires_grad True optimizer optim.SGD( model.parameters(), lr0.0001, momentum0.9, weight_decay5e-4 )解冻后学习率一定要降我习惯降到原来的十分之一。因为在ImageNet上学到的权重已经比较稳定过大的学习率会把它们从最优区域踢出去。有的项目还会对features和classifier设置不同学习率但这个技巧对AlexNet这种小网络收益不大最简单的方式还是先冻结后解冻。6.3 一个完整的微调示例假设你有一个100类的小数据集完整微调流程大概是这样model models.alexnet(weightsmodels.AlexNet_Weights.IMAGENET1K_V1) model.classifier[6] nn.Linear(4096, 100) # 第一阶段冻结features for param in model.features.parameters(): param.requires_grad False optimizer optim.SGD(model.classifier.parameters(), lr0.001, momentum0.9) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) # 跑5个epoch左右观察val_acc变化 # 之后解冻lr降到1e-4再跑20个epoch有个容易忽略的点冻结的层仍然会参与forward和反向传播的计算图只是参数不会更新。它不会帮你省多少显存因为前向还是要算。真要从根本上省显存得用小batch size或更小的输入尺寸。7. 实操踩坑记录与性能建议7.1 显存溢出不是模型错是batch太大我刚开始在8G显卡上跑batch_size128224x224的AlexNet加载完数据直接OOM。你以为需要换显卡其实只需要调小batch_size。AlexNet的显存占用大头在全连接层4096维的中间特征在反向传播时会占不少显存。把batch_size从128降到64甚至32一切正常。如果你想保持大batch可以试试torch.cuda.amp混合精度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()混合精度在NVIDIA显卡上能明显提速并减少显存代价是偶尔出现数值不稳定不过现在PyTorch的amp已经相当成熟损失函数和优化器都不用改。初次调试建议先关amp跑通后再开。7.2 LRN该不该删一个工程上的权衡LRN在PyTorch里的实现是nn.LocalResponseNorm我建议你做一个对比实验保留LRN vs 替换成BatchNorm。在我的数据集上LRN确实会让训练速度慢10%左右而最终准确率几乎没有差别。所以如果项目目标是快速验证自己的想法直接删掉LRN或换成BatchNorm。但如果你是在复现论文数值或者做算法对比实验那就保留因为改动一个模块会影响baseline的一致性。7.3 训练不收敛的排查顺序不少读者说“我照着代码改了但loss不降”我一般按这个顺序排查数据预处理确认图像没有整成0到1之外的异常值Normalize的mean/std是否用错。标签确认ImageFolder生成的标签从0开始且连续最后一个类别不会比num_classes大。模型输出确认最后一层输出维度等于类别数而不是等于1000就忘了改。学习率从头训练时0.01太大可以先试0.001Adam试0.0001。损失如果loss在训练一开始就跳到NaN优先检查数据和lr而不是模型结构。这五步能解决90%的“不收敛”。还有一个小技巧先拿一个batch过一遍模型打印输出shape再算loss确保前向和损失函数是通的。7.4 在不同硬件上的速度表现我没有专业测试环境只能给出个人经验值。在RTX 3060级别显卡上batch_size64约1.3万张图的训练集从头训练AlexNet一个epoch大约几分钟在CPU上则是时间灾难一小时能跑完一个epoch都算快。如果只有CPU建议把输入尺寸降到160x160试一下AlexNet结构不变只是先把验证流程跑通。场景batch_size设备单epoch耗时参考1.3万张训练图64RTX 3060约2~4分钟同上16CPU多核约20~40分钟微调冻结features64RTX 3060约1~2分钟这些数字仅供量级参考实际跟数据读取速度、线程数、显卡调度都有关系。关键是别用CPU做完整训练那是浪费生命。8. 写在最后的个人体会说实话我从零手写AlexNet已经是很多年前的事了但那一次经历让我把卷积、感受野、Dropout、LRN这些东西从“听说过”变成“真懂”。现在做项目我依然经常把它当作第一个baseline结构简单调参容易训练速度快出错了也好排查。如果你想真正掌握PyTorch的完整训练流程与其一开始就上几十层的残差网络不如先把AlexNet完整跑通一遍。这篇博文里的代码和踩坑经验都是我实际用过的你复制下来改改路径就能跑。以后遇到数据量更小、类别更复杂的情况你自然会知道哪些地方该调、哪些地方不该动。祝你的模型早日收敛。