3、PyTorch模型基础：从零构建一个简单的CNN分类模型（ResNet-18）

说实话，很多搞量化部署的朋友，一上来就盯着算子替换、精度校准这些高级话题。但我个人觉得，如果连模型本身的结构都没吃透，后面做量化优化很容易踩坑。今天我们就从零开始，手撸一个ResNet-18，把模型定义、前向传播、参数存储这三块彻底搞明白。

3.1 为什么选ResNet-18？

你可能会问：讲CNN分类模型，为啥不选VGG或者AlexNet？嗯，这里有个关键原因。我在做GPU推理加速时发现，ResNet的残差结构对量化部署特别友好。为什么呢？因为残差连接能缓解梯度消失，让深层网络也能稳定训练。而且ResNet-18结构相对规整，层数不多不少，非常适合作为我们后续做量化实验的基准模型。

核心要点：ResNet-18由4个残差块组（Layer1~Layer4）组成，每个组包含2个BasicBlock，总共8个BasicBlock + 1个初始卷积层 + 1个全连接层。

3.2 模型定义：从BasicBlock开始

我们先定义最基础的组件——BasicBlock。说白了，这就是一个带有残差连接的小模块。

import torch
import torch.nn as nn
import torch.nn.functional as F

class BasicBlock(nn.Module):
    expansion = 1  # 输出通道扩展系数

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, 
                               stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, 
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample  # 用于匹配维度

    def forward(self, x):
        identity = x  # 保存输入，用于残差连接

        out = self.conv1(x)
        out = self.bn1(out)
        out = F.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        # 如果输入输出维度不匹配，需要下采样
        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity  # 残差相加
        out = F.relu(out)

        return out

注意看，这里有个downsample参数。我在项目中遇到过，很多人写ResNet时忘记处理维度不匹配的问题，导致模型跑起来就报错。当输入输出通道数不同，或者特征图尺寸减半时，我们需要用1x1卷积把输入投影到正确的维度。

3.3 组装ResNet-18

有了BasicBlock，我们就可以像搭积木一样组装ResNet-18了。

class ResNet18(nn.Module):
    def __init__(self, num_classes=1000):
        super(ResNet18, self).__init__()
        self.in_channels = 64

        # 初始卷积层
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # 4个残差块组
        self.layer1 = self._make_layer(64, 2, stride=1)
        self.layer2 = self._make_layer(128, 2, stride=2)
        self.layer3 = self._make_layer(256, 2, stride=2)
        self.layer4 = self._make_layer(512, 2, stride=2)

        # 分类头
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

    def _make_layer(self, out_channels, blocks, stride):
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels, kernel_size=1, 
                          stride=stride, bias=False),
                nn.BatchNorm2d(out_channels),
            )

        layers = []
        layers.append(BasicBlock(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for _ in range(1, blocks):
            layers.append(BasicBlock(out_channels, out_channels))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

个人经验：我建议你在写_make_layer时，把self.in_channels的更新逻辑放在循环外面。曾经有个同事把更新写在了循环里，结果每个BasicBlock的输入通道都变成了上一层的输出通道，模型直接崩了。

3.4 前向传播：数据是怎么流动的？

我们来看一张图，直观理解数据在ResNet-18中的流动路径。

ResNet-18 前向传播数据流图 输入图像 (3x224x224) Conv1 (64x112x112) MaxPool (64x56x56) Layer1 (64x56x56) Layer2 (128x28x28) Layer3 (256x14x14) Layer4 (512x7x7)

你看，数据从输入开始，经过初始卷积和池化，尺寸从224x224缩小到56x56。然后经过4个残差块组，每经过一组，通道数翻倍、尺寸减半。最后通过全局平均池化变成1x1的向量，送入全连接层分类。

注意：在Layer2、Layer3、Layer4的第一个BasicBlock中，stride=2会导致特征图尺寸减半。这时候残差连接必须通过downsample来匹配维度。我曾经见过有人忘记设置downsample，结果模型训练时loss直接变成NaN。

3.5 参数存储：模型到底存了什么？

模型定义好了，参数是怎么存的？说白了，PyTorch的state_dict就是一个有序字典，里面存了所有可训练参数的名称和数值。

# 实例化模型
model = ResNet18(num_classes=10)

# 查看模型结构
print(model)

# 查看参数统计
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"总参数量: {total_params:,}")
print(f"可训练参数量: {trainable_params:,}")

# 查看state_dict的键
state_dict = model.state_dict()
print(f"\nstate_dict包含 {len(state_dict)} 个键值对")
for key in list(state_dict.keys())[:5]:  # 只显示前5个
    print(f"  {key}: {state_dict[key].shape}")

输出结果大致如下：

总参数量: 11,173,962
可训练参数量: 11,173,962

state_dict包含 62 个键值对
  conv1.weight: torch.Size([64, 3, 7, 7])
  bn1.weight: torch.Size([64])
  bn1.bias: torch.Size([64])
  layer1.0.conv1.weight: torch.Size([64, 64, 3, 3])
  layer1.0.bn1.weight: torch.Size([64])
组件 参数量 占比
初始卷积层 (conv1) 9,408 0.08%
Layer1 (2个BasicBlock) 147,456 1.32%
Layer2 (2个BasicBlock) 524,288 4.69%
Layer3 (2个BasicBlock) 2,097,152 18.77%
Layer4 (2个BasicBlock) 8,388,608 75.08%
全连接层 (fc) 5,130 0.05%

关键发现：你看，Layer4占了75%的参数量！这意味着什么？在做量化部署时，Layer4的权重对精度影响最大。我一般会优先对Layer4做更精细的量化策略，比如逐通道量化或者保留更高位宽。

3.6 保存与加载模型

模型训练好了，怎么保存？我推荐只保存state_dict，而不是整个模型对象。为什么呢？因为只保存参数更灵活，加载时你可以先定义模型结构，再加载参数，这样代码更清晰。

# 保存模型参数
torch.save(model.state_dict(), 'resnet18_cifar10.pth')

# 加载模型参数
model = ResNet18(num_classes=10)
model.load_state_dict(torch.load('resnet18_cifar10.pth'))
model.eval()  # 切换到推理模式

避坑指南：我曾经在加载模型时忘记调用model.eval()，结果推理结果一直不对。因为BN层和Dropout在训练和推理模式下行为不同。记住：推理前一定要调eval()，否则你会被坑得很惨。

3.7 小结

今天我们做了三件事：定义了ResNet-18的完整结构，理解了数据前向传播的路径，搞清楚了参数是怎么存储的。这些基础打牢了，后面做量化优化时你才能知道：哪些层对精度敏感，哪些层可以大胆压缩。

嗯，我个人觉得，花时间把模型结构吃透，比盲目调参要划算得多。你想想看，如果连模型有多少层、每层什么作用都不清楚，后面做量化部署时出了问题，你连排查的方向都没有。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →