☰
PyTorch实现经典LeNet:从参数量计算到模型保存加载
2026/10/14 8:38:34 网站建设 项目流程

我做了三年多的深度学习落地项目,有一个特别真切的感受:很多看着"老掉牙"的模型,反而是最好的老师。LeNet-5就是这么一个典型——1998年提出的网络,到现在依然是无数人入门卷积神经网络的第一个完整模型。PyTorch里跑通LeNet很简单,几十行代码的事,但真正把"结构设计、维度推导、参数量计算、模型持久化"这一整套流程理清楚,你才算真的吃透了它。这篇文章就围绕LeNet在PyTorch下的完整实现展开,重点拆解参数量是怎么算出来的,以及模型保存和加载里那些文档不会明说的坑。不管你是刚接触深度学习的学生,还是想快速把经典网络落地到业务里的工程师,这套思路都值得过一遍。

1. LeNet-5的结构到底在教我们什么

1.1 一个二十多年前的结构为何至今仍值得手写

LeNet-5(我在后文统称LeNet)是早期卷积神经网络的中流砥柱,它定义了卷积网络的基本范式:卷积层提取特征、下采样降低分辨率、全连接层做分类。因为结构小、训练快、依赖少,它成了无数人接触CNN的第一站。写LeNet最大的价值在于,你能在完全不依赖预训练权重的情况下,亲手把每一层的输入输出尺寸推一遍,体会卷积核、步长、填充这些概念对张量形状的真实影响。

LeNet的经典结构接收 1×32×32 的灰度输入,经过两层卷积(配合两次池化),再接三层全连接输出。和现在动辄上百层的ResNet、Transformer比,它就只有几个部件,但每个部件都不可或缺。我经常跟团队里新人说一句话:先把LeNet的每一行代码手写明白,后面再上大模型心里才有底。这不是情怀,是性价比最高的学习路径。

值得注意的一点是:LeNet的原始论文输入是32×32,MNIST数据集默认是28×28。很多人直接拿MNIST训练,不改padding或结构调整,模型照样能跑,但输出维度和论文对不上,后续参数量计算也会出错。所以我后面给出的代码会明确按 32×32 输入来处理。

1.2 每一层的张量维度变化:用"尺子"量一遍

在写代码前,我强烈建议你用一张纸把每一层输出尺寸算一遍。卷积层的通用计算公式是:

[ 输出高宽 = \frac{输入高宽 + 2 \times padding - kernel_size}{stride} + 1 ]

LeNet的C1层:输入 1×32×32,卷积核 5×5,输出通道6。padding=0,stride=1,所以输出高宽是 (32 - 5) + 1 = 28,得到 6×28×28。S2层是平均池化,窗口2×2、步长2,输出直接减半到 6×14×14。C3层:输入 6×14×14,用 5×5 卷积输出16通道,尺寸变为 (14 - 5) + 1 = 10,得到 16×10×10。S4层再池化到 16×5×5。从这之后进入全连接部分,把 16×5×5 展平成400维向量,依次经过120、84两个全连接层,最后输出10个类别。

这个推导过程极其重要。每一层的输出尺寸都决定了下一层的输入维度,任何一个地方算错,后面的全连接层维度就得跟着改。我第一次带着某位实习生做这个网络时,他跳过了手算直接写代码,结果全连接层的输入维度反复报错。后来我让他把上面这张"尺子"画出来,再去看PyTorch的报错信息,就立刻明白问题出在哪了。

下表是我整理的维度变化汇总,建议你对着自己画的版本核对一遍:

层名操作输入形状输出形状关键参数
C1卷积1×32×326×28×28kernel=5, stride=1, pad=0
S2平均池化6×28×286×14×14kernel=2, stride=2
C3卷积6×14×1416×10×10kernel=5, stride=1, pad=0
S4平均池化16×10×1016×5×5kernel=2, stride=2
C5卷积16×5×5120×1×1kernel=5, stride=1, pad=0
F6全连接12084ReLU激活后输出
OUTPUT全连接8410分类输出

C5这个层在不同实现里有两种处理方式:一种是像我这样写成卷积层(kernel_size=5,输出恰好是1×1),另一种是直接改成Flatten后接全连接层。两种写法参数量完全一样,但用卷积写更贴近论文原意,也方便后面对比参数量。

1.3 一个容易被忽略的细节:输入预处理

LeNet的设计输入是单通道灰度图。如果你用的是彩色图片,比如三通道RGB,第一层卷积的in_channels就要改成3。这个改动会直接影响参数量——C1层的权重数量从 1×6×5×5 变成 3×6×5×5。很多项目里模型效果不好,排查半天最后发现是通道数对不上。

我在自己的项目里通常会在数据加载阶段统一做灰度化和Resize到32×32,而不是在模型内部改结构。这样模型保持LeNet的原始定义,通用性更强,换数据集时不用改网络。归一化用 mean=0.5、std=0.5 就够,MNIST这类简单数据集不需要太重的前处理。

2. PyTorch实现LeNet:逐行写清楚

2.1 网络定义:用 Sequential 还是手写 forward

PyTorch里写LeNet主要有两种风格:全部用nn.Sequential堆叠,或者在__init__里声明模块后手写forward。我推荐后者,原因有三:一是中间层的输出在调试时更容易打印;二是后面讲参数量计算时,每个模块有独立的名字,访问权重更方便;三是和论文结构一一对应,别人读你的代码时不用猜。

下面是按论文结构实现的代码,输入是 1×32×32,输出10类:

import torch import torch.nn as nn class LeNet(nn.Module): def __init__(self, num_classes=10): super(LeNet, self).__init__() # 卷积特征提取部分 self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0) self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0) self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2) self.conv3 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0) # 全连接分类部分 self.fc1 = nn.Linear(in_features=120, out_features=84) self.fc2 = nn.Linear(in_features=84, out_features=num_classes) def forward(self, x): x = self.pool1(torch.relu(self.conv1(x))) x = self.pool2(torch.relu(self.conv2(x))) x = torch.relu(self.conv3(x)) x = x.view(x.size(0), -1) # 展平:batch_size × 120 x = torch.relu(self.fc1(x)) x = self.fc2(x) return x

注意这里卷积层后面跟的是ReLU而不是论文原始的tanh激活。这不是偷懒,而是现代训练的常规做法:ReLU计算更简单、梯度更平稳,在MNIST这种任务上收敛更快。你用 tanh 也能跑,但实测收敛速度会慢一些,初学者容易误以为模型出了问题。

2.2 卷积的3D特征图与全连接的展平逻辑

很多新手对x.view(x.size(0), -1)这个操作不理解,总觉得自己少学了一个概念。其实它做的事非常简单:把一张 120×1×1 的特征图拉成一个120维的向量。x.size(0)是batch维度,保留它;-1表示由PyTorch自动推断,把剩下的所有元素压平。

这里必须强调一个容易踩的坑:不要写成x.view(-1, 120),因为你会丢掉batch维度的第一维。如果是单样本推理,两者结果看起来一样,但一进入batch训练,模型全连接层的输入维度就全乱了。我见过不止一个项目里的代码在view这行翻车,报错信息五花八门,根子都是张量形状和预期不一致。

如果你希望代码更"现代化"一点,可以把x.view(x.size(0), -1)换成torch.flatten(x, start_dim=1),语义更清晰。两者效果完全相同,看个人习惯。

2.3 训练一个epoch:完整可运行的代码骨架

只有网络定义还不够,实际跑起来还需要数据加载、损失函数、优化器和训练循环。这里给一个最精简但结构完整的训练骨架:

import torch import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset = torchvision.datasets.MNIST( root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) model = LeNet(num_classes=10) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss = 0.0 for images, labels in loader: optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader) if __name__ == '__main__': for epoch in range(1, 4): avg_loss = train_one_epoch(model, train_loader, criterion, optimizer) print(f'Epoch {epoch:02d}, Loss: {avg_loss:.4f}')

跑三个epoch后,MNIST测试集上的准确率一般在97%以上。如果低于这个水平,优先检查数据预处理是否做了32×32的Resize,以及是否把图像归一化到了[-1,1]区间。很多人忽略了Resize这一步,输入变成了28×28,第一层卷积输出变成24×24,全连接的输入就不是120了,模型直接报错。

我建议你在forward里临时加一行print(x.shape)来检查每层输出形状,确认和手算的维度一致后再删掉。这种"打印shape验证推导"的习惯,比任何调试工具都管用,尤其是刚上手的时候。

3. 计算模型参数量:三种方法,从手算到API

3.1 为什么参数量不是"模型有多少层"这么简单

参数量就是模型里所有可训练张量的元素数量之和,它包括权重(weight)和偏置(bias)。对LeNet来说,主体贡献来自卷积层和全连接层的权重矩阵。激活函数、池化层和展平操作因为不包含可训练参数,不参与统计。

算参数量是验证你网络结构定义是否合理的重要指标。比如你以为自己的模型是个小网络,结果参数量和别人的ResNet18一个量级,那说明层设计一定有问题。反过来,如果参数量per层手算和代码统计对不上,那要么是代码定义错了,要么是手算漏了什么参数。

参数量直接影响部署时的内存占用、带宽和计算量基准。一个粗糙的估算方法是:模型参数为P,在FP32精度下,权重占4P字节;如果优化器带Adam的动量项,训练时内存开销还要再翻几倍。所以算清楚参数量不是一个可有可无的数字游戏,它是做资源规划的第一步。

3.2 手算全过程:跟着我一步步盘算盘算

手算的规则其实很简单:卷积层的权重是 [输出通道, 输入通道, 卷积核高, 卷积核宽],池化层没有参数,全连接层的权重是 [输出维度, 输入维度]。以我们上面的LeNet定义为例:

  • C1层:6个输出通道、1个输入通道、5×5卷积核。权重 = 6 × 1 × 5 × 5 = 150,偏置 = 6,合计156。
  • C3层:16个输出通道、6个输入通道、5×5卷积核。权重 = 16 × 6 × 5 × 5 = 2400,偏置 = 16,合计2416。
  • C5层:120个输出通道、16个输入通道、5×5卷积核。权重 = 120 × 16 × 5 × 5 = 48000,偏置 = 120,合计48120。
  • F6层:84个输出、120个输入。权重 = 84 × 120 = 10080,偏置 = 84,合计10164。
  • OUTPUT层:10个输出、84个输入。权重 = 10 × 84 = 840,偏置 = 10,合计850。

所有层相加:156 + 2416 + 48120 + 10164 + 850 =61706个参数。这是手算结果,后面我会用代码验证。

这里要特别提醒:池化层(AvgPool2d)是真的没有任何参数的,很多资料里会把池化也算成参数,这是错的。池化只是对特征图做固定窗口操作,没有可学习的权重。

3.3 用PyTorch的API统计:一行代码验证

PyTorch里统计参数量的标准做法是通过model.parameters()遍历所有张量,数一数每个张量有多少个元素。代码只有几行:

def count_parameters(model): total = sum(p.numel() for p in model.parameters()) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) return total, trainable model = LeNet(num_classes=10) total_params, trainable_params = count_parameters(model) print(f'Total parameters: {total_params}') print(f'Trainable parameters: {trainable_params}')

输出结果和手算的一致:总共61706个参数,其中全部都是可训练的。如果你在__init__里手动把某些层设置了requires_grad=False,trainable就会比total小,这个方法能帮你确认冻结层是否生效。

还有一个更细的层次:如果你想知道单层参数量,可以这样:

for name, param in model.named_parameters(): print(f'{name}: {param.numel()}')

输出会显示conv1.weight: 150、conv1.bias: 6等,和手算逐条对上。这种逐层对比是我排查网络定义错误时最喜欢用的方式:打印出来,手算表放旁边,哪里对不上一目了然。

3.4 更灵活的统计库:torchsummary 与 thop

如果你不想自己写统计代码,社区里有两个现成工具很好用。torchsummary可以同时显示每一层的输出形状和参数量:

from torchsummary import summary model = LeNet(num_classes=10).to('cuda' if torch.cuda.is_available() else 'cpu') summary(model, input_size=(1, 32, 32))

它会输出类似下面的信息(格式视版本略有差异):

---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 6, 28, 28] 156 AvgPool2d-2 [-1, 6, 14, 14] 0 Conv2d-3 [-1, 16, 10, 10] 2,416 AvgPool2d-4 [-1, 16, 5, 5] 0 Conv2d-5 [-1, 120, 1, 1] 48,120 Linear-6 [-1, 84] 10,164 Linear-7 [-1, 10] 850 ================================================================ Total params: 61,706

注意这里的Output Shape第一维是-1,表示batch大小。torchsummary在安装时有个小坑:最新的PyTorch版本需要安装对应的适配版本,否则会报module 'torch' has no attribute 'FloatTensor'之类的错误。要是遇到这种问题,换个环境或改用thop即可。

thop主打的是FLOPs和参数量一起统计,适合做模型轻量化对比时使用:

from thop import profile model = LeNet(num_classes=10) input_tensor = torch.randn(1, 1, 32, 32) flops, params = profile(model, inputs=(input_tensor,)) print(f'FLOPs: {flops / 1e6:.2f}M, Params: {params / 1e3:.2f}K')

但thop的FLOPs统计结果在不同PyTorch版本里会有细微差异,我建议只把它当参考值,不要当作跨版本比较的唯一依据。真要严谨,还是以手算推导和model.parameters()为准。

4. 模型保存与加载:不止是 torch.save 这么简单

4.1 保存的两种方式:整个模型 vs 状态字典

PyTorch保存模型有两种主流写法,很多人一开始搞不清楚区别:

# 方式一:保存整个模型(不推荐) torch.save(model, 'model_complete.pt') # 方式二:保存state_dict(推荐) torch.save(model.state_dict(), 'model_state.pt')

第一种方式把模型结构、权重、甚至优化器状态都序列化进一个文件里,但它和特定类定义的耦合非常深。换个环境、改一下类的模块路径,torch.load就可能直接报错。第二种方式只保存权重张量,是一个纯粹的键值映射,和模型类定义解耦,加载时只要你用相同的模型结构初始化,然后load_state_dict就能完美恢复。

我在实际项目里保存的是state_dict,而且是按下面这种格式组织的:

checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, 'config': {'num_classes': 10} } torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pt')

把epoch、权重、优化器状态和超参数配置打包在一起,训练中断后可以无缝恢复。这在长时间训练里几乎是必须的,不是锦上添花。

4.2 加载模型时最常见的两类报错

加载模型最经典的报错是size mismatch。比如你保存模型时有10个输出类,加载时初始化成了LeNet(num_classes=100),那最后一层全连接的权重尺寸就对应不上:

RuntimeError: Error(s) in loading state_dict for LeNet: size mismatch for fc2.weight: copying a param with shape torch.Size([10, 84]) from checkpoint, the shape in current model is torch.Size([100, 84]).

这种报错的原因非常明确,但报错信息里包含是你改了模型类别数还是加载了别人的预训练权重,就需要你自己判断。我以前调试过几次后发现,最快的修复办法是:检查你想加载的checkpoint里保存的config键,看看当时用的是多少类别。没有这个键的话,就用torch.load打印出所有键名,逐一核对每一层的形状。

另一类常见错误是ModuleAttributeError,报错说模型没有某个属性。这通常是因为你定义模型时改了层名,比如把self.fc1改成了self.classifier1。state_dict里的键是带层名的,比如fc1.weight,重新命名后键对不上,加载自然失败。解决办法是:要么保持层名不变,要么加载时做一个键名映射,我通常建议前者,代码可读性也更好。

4.3 设备迁移:map_location 的正确打开方式

训练脚本最常见的一个坑是:在GPU上保存的模型,拿到没有GPU的机器上加载,直接报错。准确说是你在GPU上训练,保存了state_dict,然后在CPU环境下torch.load这个文件会得到一个带cuda设备标签的张量。

正确做法是加载时用map_location参数显式指定设备:

# 加载到CPU model_state = torch.load('model_state.pt', map_location='cpu') model.load_state_dict(model_state) # 或在GPU上加载时显式指定cuda model_state = torch.load('model_state.pt', map_location='cuda')

map_location='cpu'是个万能保险,不管原模型是在哪个设备上训练的,都能把张量搬到CPU。之后如果训练在GPU上,再用model.to('cuda')手动搬一次就行。反过来,如果你在CPU上训练好模型想搬到GPU推理,PyTorch加载后还需要手动model = model.to('cuda')。

还有一个隐藏细节:torch.load在较新版本的PyTorch里默认启用了weights_only=True的模式,这意味着它只会加载张量、字典、列表等安全类型,不会执行任意Python对象的反序列化。如果你的checkpoint里除了张量还放了自定义对象,加载时可能会失败或要求你显式设置weights_only=False。安全考虑,我建议保持默认并尽量让checkpoint里只存基础数据类型。

4.4 自定义类的反序列化问题

当checkpoint里保存的是整个模型,而不是state_dict时,torch.load会尝试重建模型实例。这个过程中PyTorch会通过pickle反序列化,这时候模型类的定义必须能被当前环境import到,并且类的模块路径要和保存时一致。一旦你重构过代码目录,比如把模型从models.lenet移到了network.arch,那加载就失败。

解决办法有两个方向。一是所有长期保存的checkpoint统一只存state_dict,避免整个模型的pickle机制;二是非要保存完整模型,就把模型类定义尽量放在一个不会变动的顶级模块里。我在团队协作项目里一直强调前者——因为同事之间代码结构各不相同,谁都说不准谁会把文件移到哪里。

如果你加载不到原始的类定义,也可以在加载后重新构造一个相同结构的模型,再手动把state_dict喂进去:

model = LeNet(num_classes=10) state = torch.load('whatever.pt', map_location='cpu') if 'model_state_dict' in state: model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state)

这样完整模型文件的序号问题都被避开了,只依赖权重的键名映射。这套思路可以应对绝大多数模型加载场景。

5. 我把这些环节串成一个完整Demo的实操记录

5.1 从定义到保存:一个端到端的脚本

把这些内容串起来,我做一个可复现的完整脚本。跟着跑一遍,你会看到从网络定义、参数量统计、训练到保存加载的完整流程:

import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader # 1. 定义模型(和前面一致) class LeNet(nn.Module): def __init__(self, num_classes=10): super().__init__() self.conv1 = nn.Conv2d(1, 6, 5) self.pool1 = nn.AvgPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.pool2 = nn.AvgPool2d(2, 2) self.conv3 = nn.Conv2d(16, 120, 5) self.fc1 = nn.Linear(120, 84) self.fc2 = nn.Linear(84, num_classes) def forward(self, x): x = self.pool1(torch.relu(self.conv1(x))) x = self.pool2(torch.relu(self.conv2(x))) x = torch.relu(self.conv3(x)) x = torch.flatten(x, 1) x = torch.relu(self.fc1(x)) x = self.fc2(x) return x # 2. 数据准备 transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_set = torchvision.datasets.MNIST( root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_set, batch_size=64, shuffle=True) # 3. 初始化与配置 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = LeNet(num_classes=10).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) # 4. 打印参数量验证 total_params = sum(p.numel() for p in model.parameters()) print(f'参数量:{total_params}') # 5. 训练两个epoch for epoch in range(1, 3): model.train() running_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch {epoch}: loss = {running_loss/len(train_loader):.4f}') # 6. 保存checkpoint torch.save({ 'epoch': 2, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': running_loss / len(train_loader), }, 'lenet_mnist_checkpoint.pt') # 7. 加载checkpoint checkpoint = torch.load('lenet_mnist_checkpoint.pt', map_location=device) model = LeNet(num_classes=10).to(device) model.load_state_dict(checkpoint['model_state_dict']) model.eval() print('模型加载完成,准备评估或推理')

代码里的model.eval()一定要加。它影响BatchNorm和Dropout这类层的行为,LeNet没有这两个层所以影响不明显,但养成习惯很重要——很多人在这个细节上吃过亏。

5.2 推理时验证"模型确实被正确加载"

保存和加载是否成功,不能只看不报错,更要有量化验证。最直接的办法是在加载前后各跑一次相同推理,对比输出是否一致:

model.eval() sample_image, _ = next(iter(torch.utils.data.DataLoader(train_set, batch_size=8))) with torch.no_grad(): output_before = model(sample_image.to(device)) # 重新加载 checkpoint = torch.load('lenet_mnist_checkpoint.pt', map_location=device) model_loaded = LeNet(num_classes=10).to(device) model_loaded.load_state_dict(checkpoint['model_state_dict']) model_loaded.eval() with torch.no_grad(): output_after = model_loaded(sample_image.to(device)) print('输出是否一致:', torch.allclose(output_before, output_after))

如果你看到True,说明权重完整恢复。如果False,基本可以断定保存和加载之间存在权重不一致或设备精度问题。输出张量因为是用同样的权重和同样的输入,理论上应该逐元素完全一致。

还有一点要留意:加载后模型应该处在评估模式,不要直接开始训练,否则BN统计量会被污染。LeNet没有BN层,但这个习惯要养成。

5.3 踩过的坑:参数统计和下采样方式的"选择"

最后分享几个使用过程中最常踩的坑,希望你能绕开。

第一,参数量统计时容易混淆"池化层"。LeNet里用的AvgPool2d没有参数,但如果你用了MaxPool2d,它同样也没有参数。真正有参数的只有卷积层和全连接层。我见过有人把池化的kernel_size也算进参数量,这是理解偏差。

第二,下采样方式会影响参数量分析与训练效果之间的对应关系。LeNet原论文用的是平均池化,但如果你换成MaxPooling,发现模型收敛更快,那也很正常。MaxPool作为下采样器在MNIST上表现更好,但这个改动不会改变参数量,因为池化没有参数。你在比较参数量时不必为此纠结,但在对比不同模型效果时要把池化方式列为变量之一。

第三,加载模型后忘了model.eval()。这个问题在没有BN和Dropout的LeNet上表现不明显,甚至准确率也不会变,但一旦换到带Dropout的模型,推理结果会随机波动。我的习惯是:只要进入推理阶段,第一行永远是model.eval()。

第四,保存路径和文件名最好不要带中文和空格,尽量用纯英文。这不光是为了跨平台兼容,在容器化部署时也省去很多转义麻烦。文件名里我一般会带上epoch和loss值,比如lenet_epoch2_loss0.08.pt,便于后续挑选checkpoint。

6. 关于参数量、保存加载的更深一层思考

6.1 参数量不是越小越好,但要"算得明白"

很多项目在优化模型时有一个误区:一味追求小参数量,好像参数越少越高级。实际部署时,参数量影响的是内存占用和权重文件大小,而真正影响推理速度的往往是FLOPs和访存带宽。LeNet参数量只有6万左右,在MNIST上跑CPU推理一毫秒都不到,参数量本身不是瓶颈。

但"算得明白"意义很大。比如你在做一个移动端部署,模型文件大小直接和参数量相关:FP32下一个参数占4字节,61706个参数就是约240KB,如果转成FP16直接减半。我会提前预估这些数据,才能判断模型能不能塞进目标设备的存储空间。手算和API统计两种方式相互印证,比单靠一个工具输出更有把握。

6.2 保存的粒度决定迭代效率

项目早期我常常图省事,每轮调参只保存最新的model。后来发现一次实验要跑十几个小时,中途断电、显存溢出、后台进程被杀,损失惨重。从那以后我开始按"最优模型+最近checkpoint"双轨保存:每轮epoch结束,如果验证集效果创新高就存一个best_model.pt,另外每N轮存一个带优化器状态的checkpoint,方便从断点恢复。

这个习惯帮我解决过很多实际问题。有一次训练LeNet的变体时发现到第12个epoch开始过拟合,由于我保存了第10轮的checkpoint,回滚只用了一分钟。没有checkpoint就得从头再来。模型保存和加载不是最后收尾时才考虑的事,而是训练过程中随时要用的保险机制。

6.3 从LeNet延伸:这套方法适用于任何网络

这篇文章虽然以LeNet为例,但参数量计算、模型保存加载的逻辑完全适用于其他模型。无论你之后写ResNet、VGG还是Transformer,五件事始终是通用的:确认每层输入输出维度、手算参数量、用API验证、保存state_dict、加载时用map_location。学会这套方法论,再看任何PyTorch模型都不会怵。

LeNet小,所以每个环节的错误暴露都很清楚。你在大模型里遇到一个shape报错,排查链会很长;在LeNet里同样的问题,5分钟就能定位。正因如此,我才一直主张把经典小网络当作练习场——在这里把所有流程跑顺,后面接触大规模模型时才能游刃有余。

我在实际使用中还发现一个小技巧:把model.named_parameters()和手算表贴到项目文档里。这样过两周回头看代码,别人甚至未来的你都能快速理解模型设计的来龙去脉。一个能给别人讲清楚参数组成的模型,通常就是一个结构性很强的模型。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询