PyTorch从零复现AlexNet：结构推导、训练调参与踩坑全记录

上手复现经典网络的时候，遇到的第一座山往往就是AlexNet。明明结构看起来不复杂，真到了自己拿PyTorch从零写一遍，卷积核大小、padding到底取多少、全连接层怎么接、训练时loss怎么死活降不下去，问题一个接一个。这篇文章就把我在PyTorch里从空文件写到能跑通训练、又能稳定收敛的完整过程拆开来讲，包括每一层的参数推导、代码怎么组织、训练推理和踩坑记录，全程会给出可以直接复制运行的完整代码块。

这篇内容适合三类人：刚学完PyTorch基础语法、想拿一个完整项目练手的人；课程作业或论文复现需要快速搞定AlexNet的人；以及已经在跑现成代码、但遇到训练不收敛或显存爆掉等异常、想搞明白根源的人。我尽量按实际动手的顺序来讲，而不是按文档目录的顺序，这样读下来你会更清楚每个环节为什么要那么写。

1. 动手前的整体思路：别急着写代码，先把结构算清楚

1.1 AlexNet四个关键创新点，以及复现时如何取舍

AlexNet之所以是里程碑，不是因为它网络层数深，而是它把当时分散的几项技术组合到了一起，并且用GPU把规模撑起来了。对我来说，复现的价值在于理解“为什么每个设计点会出现在那个位置”。

第一点是ReLU激活函数。在AlexNet之前，主流是tanh和sigmoid，但梯度在深层网络中非常容易饱和。ReLU在正区间的导数是常数1，反向传播时梯度传递更稳定，训练速度也快不少。这点在PyTorch里就是一个nn.ReLU(inplace=True)，没有实现成本，但你要理解它放在卷积和全连接之后的意义。

第二点是Dropout。AlexNet的FC层有4096维，参数量非常大，训练时很容易过拟合。Dropout在训练时随机把一部分神经元置零，迫使网络不依赖某个特定神经元，相当于在一个Batch里做了多模型的隐式集成。PyTorch里是nn.Dropout(p=0.5)，要注意的是推理时必须关了它，PyTorch的model.eval()会自动搞定，但如果你自己写推理循环，忘了切模式就麻烦了。

第三点是LRN局部响应归一化。这个在现代复现里我基本不建议用，PyTorch虽然有nn.LocalResponseNorm，但大量实验证明它对最终准确率的提升非常有限，而随着BatchNorm的出现，LRN已经被完全取代。所以我会在后面代码里用BatchNorm替代LRN的位置，效果更好，收敛也更快。

第四点是重叠池化。AlexNet用的MaxPool是3x3核、stride=2，相邻池化窗口有重叠。相比不重叠的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、stride=4、padding=2，代入公式得到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, padding=4), 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', train=True, download=True, transform=transform_train) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True) test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True)

shuffle=True只在训练集上打开，测试集不需要，因为测试时我们要逐个样本评估，打乱顺序没有任何意义。num_workers可以根据CPU核数调整，Windows上如果报错就减到0或者2。pin_memory=True在GPU训练时能加快数据从CPU内存拷贝到显存的速度，这个在数据量大时收益更明显。

2.3 数据增强策略：RandomCrop和RandomHorizontalFlip为什么够用

很多新手会纠结要不要上复杂的增强，比如色彩抖动、随机擦除、Cutout。我的建议是，在CIFAR-10上先做两个最基础的：随机裁剪和水平翻转。

RandomCrop(32, padding=4)的做法是先把32x32的图pad到40x40，再随机裁回32x32，相当于给网络提供了一些平移不变性。RandomHorizontalFlip让图片有一定概率左右翻转，对CIFAR-10里大多数物体类别都是合理变换，可以让训练数据量翻倍。

这两个增强的作用是缓解过拟合。AlexNet的参数量高达数千万，直接在小数据集上硬训，训练集loss很快会降到接近0，测试集准确率却上不去，这就是过拟合。增强相当于给网络提供了不同视角的样本，让它学习到更鲁棒的特征。

3. AlexNet网络结构逐层拆解与PyTorch代码实现

3.1 网络整体参数一览表

在写模型类之前，我把每层配置整理成了一张表，后面写代码时对照着填参数，不容易犯错。

如果把最后的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_classes=10, in_channels=3): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), ) self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) self.classifier = nn.Sequential( nn.Dropout(p=0.5), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True), nn.Dropout(p=0.5), nn.Linear(4096, 4096), nn.ReLU(inplace=True), 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, mode='fan_in', nonlinearity='relu') 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，池化后变成27；conv2保持27，池化后13；conv3、conv4、conv5都保持13，最后池化成6。如果输入换成32x32的CIFAR-10，conv1后尺寸是(32 - 11 + 4)/4 + 1 = 7，池化后3，后面几层维持3，最后池化到1，依然能通，只是特征图太小了，所以我推荐用AdaptiveAvgPool把它固定到6x6。

3.3 权重初始化为什么重要

这段代码里的_initialize_weights不是可有可无的装饰。PyTorch的nn.Conv2d默认初始化方式是Kaiming均匀分布，但那种初始化是针对没有激活函数的情况设计的。AlexNet大量使用ReLU，ReLU会把负半轴的信息直接丢弃，如果权重初始化不合理，很多神经元一开始就输出为0，梯度也一直是0，这个神经元就永久死掉了。

我采用Kaiming正太初始化，它根据ReLU的特点计算了合适的方差范围，让每一层输出的方差在传播过程中保持稳定。全连接层我用均值0、标准差0.01的正态分布，这是AlexNet原文的做法。如果注释掉初始化函数跑一轮对比，你会发现训练前期loss下降慢很多，甚至可能出现loss变成nan的情况。

4. 训练流程细节与超参数调优：照抄不翻车

4.1 损失函数与优化器配置

分类任务的标准配置是交叉熵损失。PyTorch的nn.CrossEntropyLoss已经做了两件事：先把模型最后一层的logits做softmax，再计算交叉熵。所以模型最后一层不要额外加softmax，否则会重复计算导致loss数值异常。

优化器我选SGD加动量，而不是Adam。很多新手会惯性用Adam，但在图像分类这种需要精细调学习率的场景，SGD+momentum的泛化能力通常更好。AlexNet原文用的momentum是0.9，weight_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_classes=10).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-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_size=30, gamma=0.1)

step_size=30表示每30个epoch降低一次，gamma=0.1表示乘以0.1。如果你的训练总轮数是90，那就经历了三次衰减：前30个epoch用0.01，30到60用0.001，60到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(f'Epoch [{epoch + 1}/{num_epochs}] ' f'Train 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_var；eval模式下Dropout关闭，BatchNorm用固定统计量。两个函数里都显式调用了对应模式，防止验证时还把Dropout开着导致结果忽高忽低。

4.4 训练时长与硬件选择建议

在单张RTX 3060级别显卡上，CIFAR-10训练90个epoch，batch 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_location=device) 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, dim=1) pred_class = torch.argmax(prob, dim=1).item() print(f'Predicted class: {CIFAR10_CLASSES[pred_class]}, probability: {prob[0][pred_class].item():.4f}')

map_location=device这个参数很重要。如果你在GPU上训练的模型要拿到CPU机器上推理，不写这个参数直接load会报错，报错信息通常类似Attempting to deserialize object on a CUDA device。写了map_location='cpu'能解决这个问题。

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, keepdim=True) 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_classes=10) 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 常见问题速查表

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、stride=4，如果你直接从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))把中间结果存下来，就算后边训练崩了，至少有一个接近最优的备份能继续用，这个习惯能帮你省下不少重跑时间。