手写可配置AlexNet：PyTorch自定义CNN结构与训练的完整实践

1. 为什么到了现在还要自己手动搭一个AlexNet

如果你去GitHub上搜，随便都能找到几十个现成的AlexNet实现，PyTorch官方仓库里也有torchvision.models.alexnet，一行代码就能把预训练权重加载出来。那为什么还要自己动手从零写一遍？

我最近在整理自己的模型库时，又把AlexNet手动实现了一遍。不是因为网上找不到，而是因为“自定义”这件事，恰恰是网上代码给不了你的。你会发现，等到你真的要改输入尺寸、换数据集、削减通道数、调整分类数量时，官方写死的模型类会变成一份“改起来很痛苦”的代码。自己搭一遍，弄清楚每一层的参数是怎么算出来的，哪些地方能改哪些不能改，这套经验在你后面做任何CNN模型时都会复用得上。

这篇文章更适合两类人。一类是刚学完卷积网络原理，想通过完整手写一个经典模型把知识串起来的人；另一类是有自己的数据集、想跑通完整流程——从定义网络到数据加载再到训练调参的人。我会按我自己实际走过的路子来写，代码基于PyTorch，版本在1.x以上都通用，大家可以直接复制去改。

AlexNet是2012年ImageNet竞赛的冠军，当年把分类错误率一下压到了15.3%而第二名是26.2%，是深度学习火起来的重要节点。它的结构放到现在看不算复杂，但很多现代CNN的基本套路都在里面。说句实话，现在做视觉任务很少会直接拿AlexNet当骨干网络，但作为“自定义神经网络的第一个模型”，它特别合适。因为它的层数不多，结构清晰，每一层的作用和参数都容易解释，你改任何一个地方都能直观看到效果。当你把这个模型彻底弄懂，后面再去碰ResNet、EfficientNet甚至Transformer结构，底层的工程思路都是通用的。

所谓“自定义”，我的理解不是让你改几层就完事，而是包括三个方面：网络结构可配置、数据输入可定制、训练细节可调节。这篇文章就是沿着这三条线展开的。

2. AlexNet结构拆解：每一层在干什么，哪里可以改

2.1 先把账算清楚：原版各层参数

原版AlexNet输入是227x227的RGB图像（论文里写的是224x224，实际实现中有边界处理，后面细说），一共8层带参数的层，前5层是卷积层，后3层是全连接层。我把每层的参数和输出尺寸整理成了一张表，方便后面调试时对照。

层名	类型	核尺寸/步长/填充	输出尺寸	参数量
conv1	卷积	11x11, stride=4, pad=2	96x55x55	约3.5万
pool1	最大池化	3x3, stride=2	96x27x27	0
conv2	卷积	5x5, pad=2	256x27x27	约61.4万
pool2	最大池化	3x3, stride=2	256x13x13	0
conv3	卷积	3x3, pad=1	384x13x13	约88.5万
conv4	卷积	3x3, pad=1	384x13x13	约132.5万
conv5	卷积	3x3, pad=1	256x13x13	约88.5万
pool5	最大池化	3x3, stride=2	256x6x6	0
fc6	全连接	—	4096	约3775万
fc7	全连接	—	4096	约1678万
fc8	全连接	—	1000	约409.6万

这里要特别提一下参数量。整个模型大约6000万参数，但其中卷积层总共只占了几百万，剩下超过九成都在后面三个全连接层。这一点直接影响你对“自定义”的判断：如果只是想快速验证一个小数据集，全连接层往往是最大的瓶颈，也是最值得先砍一刀的地方。

另外原版AlexNet在第一个卷积层用了11x11的大卷积核，这种大核在浅层能捕捉更多全局纹理信息，但计算量也大。现在很多轻量化网络已经把第一层改成7x7甚至3x3，配合多个卷积层叠加来获得更大的感受野。你在自定义时如果输入图像本身就很小，比如只有64x64，就不要硬套11x11核，适当调小能让网络更匹配你的数据。

2.2 卷积层堆叠的技巧：为什么是3x3和5x5混着用

AlexNet的卷积层并不是全部用同一种尺寸的核，而是11、5、3交替出现。很多人初学时不理解这个设计的用意，以为只是随便挑的。其实这里面有一个很朴素的道理：浅层用大核快速下采样并提取大范围特征，深层用小核精修局部细节，同时控制参数规模。

按stride=4、核11x11来算，输入227x227经过conv1后大约得到55x55的输出，这一步相当于把空间尺寸快速压到四分之一。到了conv3到conv5，全部换成3x3核并保持padding=1，空间尺寸就稳定在13x13不变，网络在接下来的层里专注做通道维度的特征组合。

这里有个容易被忽略的细节：卷积层的通道数增长曲线是96→256→384→384→256。它的规律是“先增后稳再降”。前两层快速增加通道数来丰富特征表达，中间两层保持最大通道数做特征融合，最后一层再降下来，目的是让信息逐步向更紧凑的表示收敛，同时也给后面的全连接层减少输入压力。你在自定义时如果想让网络更宽或更窄，建议沿用这个“增-稳-降”的通道节奏，而不是随意跳变，否则模型容易在训练初期就出现特征表达不稳定的问题。

2.3 LRN和Dropout：老结构里的防过拟合三板斧

原版AlexNet里有三个防止过拟合（或提升泛化）的手段：局部响应归一化（LRN）、重叠池化、Dropout。

LRN现在用得很少了，很多开源实现里干脆把它去掉。它的作用是在相邻通道间做归一化，模拟神经生物学里的“侧抑制”，让响应大的通道抑制周围通道。在实践中我发现，LRN对最终精度影响很小，但它增加了计算量，而且在某些框架上实现不一致，容易导致导入导出时出问题。如果你只是在自己数据集上训练，可以大胆去掉LRN，用一个BatchNorm2d来替代，效果通常会更好。

重叠池化指的是池化窗口3x3、步长2，窗口有重叠。相比现在常用的2x2池化+步长2，重叠池化保留了更多信息，也能轻微缓解过拟合。同样，这个设计在现代网络中基本被stride=2的卷积下采样取代了，你在自定义时完全可以直接用Conv2d(stride=2)来缩小特征图，效果更稳定。

Dropout是整个AlexNet里最重要的防过拟合手段之一。原版在两个全连接层后各接了一个ratio=0.5的Dropout，训练时随机屏蔽一半的神经元，推理时自动关闭。这里有一个实战经验：不要只在最后两个全连接层加Dropout，如果你把模型改成了只有一层全连接头，建议这一层后面仍然保留Dropout，p值可以从0.5降到0.3左右，防止把信息丢得太多。

2.4 从全连接到全局平均池化：自定义时最常动的一刀

AlexNet原版用了三层全连接，每层4096，这是参数爆炸的主要来源。当你只需要在一个小数据集上做分类时，这个设计是很大的浪费。我自己的习惯是：把原版的三层全连接收敛成一层或两层，或者干脆在最后的卷积特征图后面接一个全局平均池化（Global Average Pooling，GAP），再直接接分类层。

GAP的做法是对每个通道的特征图求平均值，得到一个长度为通道数的向量。这样做的好处是：不管输入图像尺寸怎么变，GAP的输出维度始终固定为C（通道数），全连接层不用跟着改输入尺寸。这给“自定义”带来了极大的灵活性。比如你训练时用的是224x224，部署时换成300x300，只要网络里的卷积层全是padding保持尺寸不变的，GAP层就能自动兼容，不会报维度错误。

所以我的建议是：如果你自定义AlexNet的主要目的是在自有数据集上做实验，第二层和第三层全连接可以先砍掉，只保留一个全连接层，前面加GAP。参数总量会从6000万掉到几百万，训练速度快很多，在小数据上反而不容易过拟合。如果你是想追求精度、手里又有足够大的数据量，那保留两层全连接更稳。

3. 用PyTorch实现一个可配置的AlexNet类

3.1 基础版本：先把结构跑通

下面是我写的核心代码。和网上很多版本不同，我把结构写成了一个可配置类，用参数控制每个阶段的通道数和全连接层宽度，而不是把所有数字写死。

import torch import torch.nn as nn class CustomAlexNet(nn.Module): def __init__(self, num_classes=10, in_channels=3, dropout=0.5, use_bn=True, use_gap=False, fc_widths=(4096,)): super().__init__() # 特征提取部分 self.features = nn.Sequential( self._conv_block(in_channels, 48, 96, 11, 4, 2, use_bn), nn.MaxPool2d(kernel_size=3, stride=2), self._conv_block(96, 128, 256, 5, 1, 2, use_bn), nn.MaxPool2d(kernel_size=3, stride=2), self._conv_block(256, 192, 384, 3, 1, 1, use_bn), self._conv_block(384, 192, 384, 3, 1, 1, use_bn), self._conv_block(384, 128, 256, 3, 1, 1, use_bn), nn.MaxPool2d(kernel_size=3, stride=2), ) # 分类部分 if use_gap: self.classifier = nn.Sequential( nn.Dropout(dropout), nn.Linear(256, num_classes) ) else: last_conv_out = 256 # 需要根据输入尺寸计算flatten后的维度 # 这里默认输入是227x227，池化后是6x6 fc_layers = [] in_features = last_conv_out * 6 * 6 for width in fc_widths: fc_layers.append(nn.Dropout(dropout)) fc_layers.append(nn.Linear(in_features, width)) fc_layers.append(nn.ReLU(inplace=True)) in_features = width fc_layers.append(nn.Dropout(dropout)) fc_layers.append(nn.Linear(in_features, num_classes)) self.classifier = nn.Sequential(*fc_layers) # 初始化 self._init_weights() def _conv_block(self, in_ch, mid_ch, out_ch, k, s, p, use_bn): layers = [nn.Conv2d(in_ch, mid_ch, k, s, p), nn.ReLU(inplace=True)] if use_bn: layers.append(nn.BatchNorm2d(mid_ch)) layers += [nn.Conv2d(mid_ch, out_ch, 1, 1, 0), nn.ReLU(inplace=True)] return nn.Sequential(*layers)

等等，我上面的代码里那个_conv_block写得有点绕，把原本的单卷积拆分成了两个卷积，这不是原版的结构，而是加了点“自定义”口味——用1x1卷积做通道变换。如果你想要最接近原版的版本，其实可以直接这样写：

def conv_block(in_ch, out_ch, k, s, p, use_bn): layers = [nn.Conv2d(in_ch, out_ch, k, s, p), nn.ReLU(inplace=True)] if use_bn: layers.append(nn.BatchNorm2d(out_ch)) return nn.Sequential(*layers)

我建议第一次实验用这个简化版，先确保网络能跑通，再去加各种“自定义花活”。1x1卷积适合在通道数特别多的时候用来降维，但在AlexNet这种结构里不是必须的。

3.2 让输入尺寸不再写死：自适应池化的技巧

上面代码里有个地方容易踩坑：如果不使用GAP，全连接层的输入维度是和输入图像尺寸强相关的。比如227x227输入经过三次池化（每次约减半）得到6x6，但如果你换成128x128输入，最后特征图就不是6x6了，Linear层的in_features就会对不上，直接报错。

解决办法有两个：要么每次按公式算特征图尺寸，要么用自适应池化。所谓自适应池化就是nn.AdaptiveAvgPool2d((6, 6))，它会把任意大小的输入特征图都缩放到6x6，这样全连接层的输入维度就固定了。这个技巧在做自定义时非常实用，因为你不一定总用224x224的输入。

我习惯在features最后面加一个AdaptiveAvgPool2d，这样既能兼容不同输入尺寸，又能在需要时随时切回GAP模式。如果你在意精度，可以只在池化层后加，不要把它当作普通的全局平均池化来理解。

3.3 参数初始化为什么重要：别用PyTorch默认值

很多人搭完模型直接开始训练，根本没想过参数初始化的问题。PyTorch的nn.Conv2d和nn.Linear默认初始化在某些任务上效果还行，但复现AlexNet这类老模型时，按原论文的方式初始化往往收敛更稳定。

原论文用的是均值为0、标准差为0.01的高斯分布初始化卷积核和全连接层权重，偏置初始化为1。这个“偏置初始化为1”是当时的一个工程技巧，目的是让ReLU在初始阶段不容易进入死区。放到今天的PyTorch里，更常见的做法是Kaiming初始化配合零偏置，效果已经足够好。

def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) 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)

如果你是偏执型选手，想完全复现原版初始化，把nn.Linear的权重标准差改成0.01、偏置设为1也可以，差别不会太大。关键是不要用默认初始化然后跳过这一步，尤其是当你自定义过网络结构、层的宽度变了之后，默认初始化的适配性会更差。

3.4 验证网络：前向传播和参数汇总

写完模型第一件事不是训练，而是先打印结构、跑一个随机输入看看输出维度对不对。

net = CustomAlexNet(num_classes=10, use_gap=True) print(net) x = torch.randn(2, 3, 227, 227) # batch=2 out = net(x) print("output shape:", out.shape) # 期望 [2, 10] total_params = sum(p.numel() for p in net.parameters()) trainable_params = sum(p.numel() for p in net.parameters() if p.requires_grad) print(f"total params: {total_params / 1e6:.2f}M, trainable: {trainable_params / 1e6:.2f}M")

这一步能提前暴露80%的维度问题。尤其是当你改了fc_widths或输入尺寸之后，务必跑一遍再进训练循环。我自己见过不少同学改完网络结构直接训练，等到第一个epoch前向传播报错才回头看，白白浪费半个小时。

还有一个小技巧：把输入尺寸写成(2, 3, 224, 224)或(1, 3, 128, 128)分别跑一次，确保自适应池化真的起作用。如果你打算以后换尺寸训练，提前测试会让你省很多事。

4. 自定义数据加载：从自己的图片文件夹到DataLoader

4.1 目录结构约定：从杂乱数据到标准格式

数据加载是整个自定义流程里最容易被低估的一环。很多数据集下载下来并不是整齐的文件夹结构，而是散落的各种图片。我在自定义实验中通常会先把数据处理成规范格式，再写Dataset去读。推荐的文件夹结构如下：

data/ ├── train/ │ ├── class1/ │ │ ├── img_001.jpg │ │ └── img_002.jpg │ ├── class2/ │ │ ├── img_001.jpg │ │ └── img_002.jpg └── val/ ├── class1/ │ ├── img_003.jpg └── class2/ └── img_004.jpg

这么做的好处有两点：一是torchvision.datasets.ImageFolder可以直接读取这种结构并自动按文件夹名字生成类别索引，省去手写CSV映射的麻烦；二是训练集和验证集分开归档，后面做划分也简单。如果你数据量不大，也可以先用脚本按比例随机分到train和val两个目录，再来训练。

4.2 自定义Dataset类的写法：解决“格式不标准”的问题

如果你的图片命名混乱、标注信息在Excel或CSV里、或者是多标签分类，那ImageFolder就不够用了，需要写自己的Dataset子类。下面是我常用的一个模板，支持从CSV读取图片路径和标签：

import os import pandas as pd from PIL import Image from torch.utils.data import Dataset class CSVDataset(Dataset): def __init__(self, csv_path, img_dir, transform=None): self.df = pd.read_csv(csv_path) self.img_dir = img_dir self.transform = transform def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img_path = os.path.join(self.img_dir, row['filename']) image = Image.open(img_path).convert('RGB') label = int(row['label']) if self.transform: image = self.transform(image) return image, label

这个类虽然简单，但覆盖了大多数自定义场景。有几个细节值得注意：

• 用Image.open后一定要convert('RGB')，否则遇到灰度图或带透明通道的PNG，后面卷积层会因为in_channels=3直接炸。
• 不用在Dataset里做数据增强，增强放在transform里，这样复用性更好，验证集和训练集可以各自使用不同的transform组合。
• 如果CSV里有错误路径，最好在__init__时做一次存在性检查，有问题尽早暴露，不要在训练到一半时才开始疯狂报错。

4.3 数据增强：在自定义任务里怎么取舍

AlexNet原论文用了随机裁剪和水平翻转，放到今天的代码里，常见的一套transform是这样：

from torchvision import transforms train_transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(227), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.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, 256)), transforms.CenterCrop(227), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

这里要特别强调一下RandomResizedCrop和CenterCrop的配合。训练时用RandomResizedCrop能把图像随机裁剪成227x227，这就同时实现了尺度变化和局部裁剪，对提升模型泛化能力帮助很大；验证时不能用随机裁剪，必须用固定的中心裁剪保证结果可复现。

关于Normalize的均值标准差，很多人会直接套用ImageNet的数值。如果你的数据集和ImageNet差异很大（比如全是显微镜图像、或者全是黑白工业零件），最好自己算一遍数据集的均值和标准差。公式很简单：统计所有像素的均值和标准差，再用这些值做标准化。PyTorch里没有现成的一行命令，可以写个脚本跑一遍样本集，很快。

4.4 DataLoader的num_workers和pin_memory：影响训练速度的关键参数

数据加载如果写得不合理，GPU会经常空转。除了batch_size之外，影响最大的是num_workers和pin_memory。

from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, # Windows上建议设为0或1 pin_memory=True, drop_last=True # 防止最后一个batch过小影响BN统计 )

num_workers是开几个子进程去读数据。Linux上开到CPU核心数的一半通常没问题，Windows上因为多进程机制不一样，经常出现卡死或内存暴涨，建议先设0或1调通代码再慢慢加。pin_memory=True会让数据加载到锁页内存，GPU训练时从CPU拷贝数据更快，前提是你的机器内存够用。drop_last=True对我来说几乎必开，尤其是用了BatchNorm的模型，因为如果一个epoch最后只剩一个样本的batch，BN统计会跳动很大，导致loss曲线出现毛刺。

5. 训练循环和超参数调优：如何判断模型真的在学

5.1 训练流程的基本框架

训练代码不需要多花哨，但要保证三个核心逻辑都完整：训练阶段更新参数、验证阶段只前向不更新、每个epoch结束保存一次模型。下面是我常用的模板：

import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net = CustomAlexNet(num_classes=10, use_gap=True).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) scheduler = CosineAnnealingLR(optimizer, T_max=30) best_acc = 0.0 for epoch in range(30): net.train() train_loss, correct, total = 0.0, 0, 0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = net(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 统计 train_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() scheduler.step() # 验证 net.eval() val_loss, val_correct, val_total = 0.0, 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = net(images) loss = criterion(outputs, labels) val_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) val_total += labels.size(0) val_correct += predicted.eq(labels).sum().item() acc = 100.0 * val_correct / val_total print(f"Epoch {epoch+1}/30 | " f"Train Loss {train_loss/total:.4f} | " f"Val Loss {val_loss/val_total:.4f} | " f"Val Acc {acc:.2f}%") if acc > best_acc: best_acc = acc torch.save(net.state_dict(), "best_alexnet.pth")

这段代码有几个地方我按实战习惯做了调整。一是用optimizer.zero_grad()而不是zero_grad(set_to_none=True)，保持最大的兼容性；二是统计loss时乘以images.size(0)，这样算的是整个epoch的loss和，而不是最后一个batch的；三是每个epoch都保存一次验证集上最好的模型，避免最后几轮过拟合时把权重存下来。

5.2 Loss、优化器、学习率怎么配

AlexNet原论文用的是SGD加momentum，这个选择放在今天依然稳。Adam虽然收敛快，但在小数据集上容易过拟合，同时收敛到精度极值时不如SGD平滑。所以我默认会用SGD：

optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) scheduler = CosineAnnealingLR(optimizer, T_max=30)

学习率怎么选，这里给出一个经验范围。如果从头训练一个小数据集（几千张图），lr=0.01是一个偏高的起点，配合余弦退火能较快找到合适的区域；如果你的任务比较难或者模型用了迁移学习，那初始学习率要更小，预训练微调通常用1e-3到1e-4才稳妥。Loss用CrossEntropyLoss就好，它是LogSoftmax加NLLLoss的组合，数值上更稳定，也避免你手动写softmax再算log出错。

5.3 Batch Size与学习率的联动：很多人忽略的事

Batch Size不只是影响显存，它还直接决定你该用多大的学习率。我的经验公式很简单：如果Batch Size翻倍，学习率大概也翻倍。比如Batch Size=32时用lr=0.01，改到64就可以试试lr=0.02。原因是batch越大，每个step的梯度估计越平稳，可以用更大的步长；batch越小，梯度噪声越大，过大的学习率容易震荡。

对小数据集，Batch Size通常设16到64之间。太小的batch会让BatchNorm的统计量不稳定，太大的batch又可能在小数据集上提前收敛到较差的平坦区域。在自定义实验里，我一般从32起步，显存够用就加到64，然后用上面那个规则同步调学习率。

5.4 判断模型“真的在学”而不是“在背答案”

训练过程中最重要的不是看训练准确率有多高，而是看训练loss和验证loss之间的距离。如果两者都在下降，说明模型在学共性特征；如果训练loss一路降到很低，验证loss却开始反弹，这是过拟合的典型信号。遇到这种情况，可以先加大Dropout的p值，或者减小weight_decay，或者提前终止训练。

还有一个容易被忽略的指标是验证准确率是否在震荡。如果验证准确率忽高忽低，可能不是模型问题，而是验证集太小，或某个batch的随机性太强。我习惯验证集至少保留几百张图，最好每个类别都有几十张，否则评价指标的方差会大到完全不可信。

6. 踩坑实录：自定义AlexNet训练中最大的几个坑

6.1 维度不匹配：全连接层永远是对不上的重灾区

这是所有自定义CNN里遇到最多的报错，形式通常是RuntimeError: size mismatch for fc.weight: copying a param with shape torch.Size([...])或者前向传播直接报mat1 and mat2 shapes cannot be multiplied。

原因基本都是全连接层的in_features和你前一层实际输出对不上。原版AlexNet从输入227x227开始，三个池化层各减半一次，最后是6x6。但只要输入尺寸变了，这个数字就变。我前面写的AdaptiveAvgPool2d((6, 6))就是为这个准备的，它能保证不管前面你输入多大，到全连接层时始终是256x6x6，这样in_features=256*6*6就永远是对的。

如果你不想用这个技巧，那就老老实实自己算：卷积层输出尺寸 = floor((输入 + 2*padding - kernel_size) / stride) + 1，池化层同理，逐层算完再填in_features。但我还是推荐自适应池化，省事且不会出错。

6.2 显存OOM：先分清是模型问题还是训练配置问题

OOM（Out of Memory）是另一个高频坑。我的排查顺序是：先用一个很小的batch，比如2，跑一遍前向和反向，看是不是模型本身太大或输入尺寸太大。能用小batch跑通，说明是batch size、输入图尺寸或数据加载的问题；如果连batch=2都爆，那大概率是模型结构需要瘦身。

针对AlexNet，最容易爆显存的位置就是全连接层，因为参数量和中间激活值都巨大。解决办法很直接：砍全连接层数、把fc_widths从(4096,)降到(512,)，或者用GAP替代全连接。也可以把输入尺寸从227降为160，显存占用会明显下降。图像分类任务里，输入尺寸从224降到160通常只会损失两三个点的精度，但显存可能减少一半还多。

6.3 学习率太高导致loss变成NaN

训练中loss突然变成NaN，新手第一反应是数据有问题，但很多时候是学习率太高，梯度过大导致的数值溢出。尤其当你自定义了网络结构后，不同层的梯度尺度可能差异很大，过大的学习率会让某一层权重更新后产生NaN，之后整个网络都跟着坏掉。

我的处理建议：如果第一次训练出现NaN，先确认数据里没有NaN（图片像素值范围异常、标签越界等）；确认数据没问题后，把学习率除以10再试。如果模型从头训练，可以考虑加一个torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)，把梯度裁剪到合理范围，防止单步更新过大。

还有一个容易被忽略的点：使用nn.CrossEntropyLoss时标签必须是0到num_classes-1之间的整数。如果CSV标签是从1开始数的，模型往往会先出现loss下降缓慢，然后突然跳到一个接近NaN的值，这也算是自定义数据加载里的经典坑。

6.4 数据量少怎么办：迁移学习和冻结层

自定义AlexNet如果从头训练而数据只有几千张，效果往往很差。这时候直接加载ImageNet预训练权重再微调是更聪明的做法。PyTorch官方提供torchvision.models.alexnet(weights=AlexNet_Weights.IMAGENET1K_V1)，你只需要把最后一层替换成自己的分类层即可。

import torchvision.models as models net = models.alexnet(weights=models.AlexNet_Weights.IMAGENET1K_V1) net.classifier[6] = nn.Linear(4096, num_classes)

这里要解释一下为什么这种微调在小数据上有效：预训练模型的前几层卷积学到的都是通用视觉特征，比如边缘、纹理、形状，这些特征在任何图像任务里都有用。只有最后一两层和分类层是跟具体任务强相关的。所以迁移学习的常规做法是：先冻结前面的卷积层，只训练新的分类层；等分类层收敛后再解冻少量深层卷积层，用小学习率微调。

如果你用的是我上面自定义的CustomAlexNet，想加载官方权重会有点麻烦，因为结构不完全一致。我的建议是：第一次实验直接用官方模型做迁移学习，等流程跑通后，再对比一下自己从头训练的自定义模型。两者的差距能给你很多直观感受。

6.5 保存和加载模型时容易踩的strict坑

PyTorch保存模型有两种常见方式：只保存权重state_dict，或者保存完整的model对象。我倾向保存state_dict，因为它只包含张量数据，跨设备迁移方便，兼容性也好。

加载时有个高频报错是Missing key(s) in state_dict和Unexpected key(s)。出现这个报错，最常见的原因是模型结构定义和state_dict不完全对应——你把use_gap改成True保存了一版，加载时又用use_gap=False的类去读，自然对不上。解决办法是：加载前先打印一下net.state_dict().keys()和保存文件的keys，一眼就能看出哪里不匹配。加载用于推理时，记得加strict=True（默认就是True）让不匹配直接报错，别用strict=False把错误悄悄吞掉，否则后面跑出来的结果莫名其妙。

7. 模型导出与后续扩展：训练完不是终点

7.1 保存checkpoint的最佳实践

训练结束后，我只保存最少的必要信息：model_state_dict、优化器状态（可选）、epoch数、验证集准确率。这样文件体积小，加载时也不需要额外的扩展库。

torch.save({ 'epoch': epoch + 1, 'model_state_dict': net.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_acc': best_acc, }, f"checkpoint_alexnet_epoch{epoch+1}.pth")

加载时需要先创建模型实例，再加载对应键：

ckpt = torch.load("checkpoint_alexnet.pth", map_location='cpu') net.load_state_dict(ckpt['model_state_dict'])

把map_location设成cpu是我反复提的建议。即使你训练时用的是GPU，在只有CPU的机器上加载时也能正常工作。如果你保存时在GPU上，加载到CPU会报device mismatch，提前加map_location='cpu'能避开这个坑。

7.2 导出ONNX做推理和部署

如果你训练完想把模型接到业务系统里，用PyTorch的torch.jit或者导出ONNX都很方便。ONNX的好处是中间表示通用，能转成各种运行时格式。导出代码很短：

net.eval() dummy_input = torch.randn(1, 3, 227, 227) torch.onnx.export( net, dummy_input, "alexnet_custom.onnx", input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} )

有几个注意点：导出前务必net.eval()，因为Dropout和BatchNorm在训练和推理时的行为不同；dynamic_axes加上之后，推理时batch大小可以不固定，对服务端部署很有用。导出后最好用onnxruntime或onnx.checker测一下，确认输出和PyTorch原模型一致。

7.3 沿着这个“自定义”的思路还能扩展什么

当你把一个可配置的AlexNet跑通之后，里面涉及的每一个点都可以向外扩展。比如把_conv_block里的普通卷积替换成深度可分离卷积，再把通道数按宽度因子缩放，这就是MobileNet那套轻量化思路的雏形；比如把fc_widths改成多个值并通过试验对比，就能理解模型容量和过拟合之间的权衡；再比如把features部分换成ResNet的BasicBlock，分类头保留当前设计，其实就迈进了“自定义ResNet”的门。

我从实际操作中得到的体会是：模型结构从来不是最重要的东西，真正重要的是“知道每一层在干什么、参数是怎么流动的、出现问题从哪里排查”这一整套工程思维。自定义AlexNet是练这套思维性价比最高的起步项目，因为它足够简单，简单到你可以把所有精力花在理解“自定义”本身，而不是被复杂结构绕晕。

自己在本地跑完一遍之后，再做三件事会有很大收获：第一件，把输入尺寸从227改成64，观察精度变化；第二件，把全连接层宽度从4096改成128，对比训练速度和显存占用；第三件，用一个很小的数据集同时从头训练和迁移学习，体会两种方式的差异。做完这三个小实验，你对神经网络“可配置”这件事的理解，会比只看十篇教程都更扎实。