3、PyTorch模型基础：torch.nn.Module、模型定义、模型保存与加载

好，咱们进入正题。这一章讲的是PyTorch模型最核心的几个概念。说白了，你后面所有TensorRT的转换工作，都建立在对这几个东西的理解上。我见过不少新手，模型定义得乱七八糟，保存加载也搞不清楚，结果到了转TensorRT那一步，各种报错，排查起来特别痛苦。

所以，咱们把地基打牢。这一章，我会带你过一遍torch.nn.Module、模型定义、以及state_dict和torch.jit.trace的保存与加载。嗯，都是干货。

3.1 torch.nn.Module：一切模型的基类

在PyTorch里，所有神经网络模型都继承自torch.nn.Module。你可以把它理解成一个乐高积木的基座。你所有的层、参数、前向传播逻辑，都得放在这个基座上。

我个人习惯，每次定义一个新模型，第一件事就是写class MyModel(nn.Module):。然后，在__init__里定义你需要的层，在forward里定义数据怎么流过这些层。

举个例子，一个简单的线性模型：

import torch
import torch.nn as nn

class SimpleLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super(SimpleLinear, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.linear(x)

你看，就这么简单。但这里有个坑，我刚开始用PyTorch时踩过——记得调用super().__init__()。不调用的话，模型内部的参数注册机制会出问题，你训练的时候梯度可能都传不回来。

注意： 在 __init__ 中定义的层，PyTorch会自动追踪它们的参数。如果你用普通的Python列表来存层，参数是不会被注册的。要用 nn.ModuleList 或 nn.Sequential。

3.2 模型定义：从简单到复杂

模型定义其实就两件事：搭积木（定义层）和连水管（定义前向传播）。

咱们看一个稍微复杂点的例子，一个两层的全连接网络：

class TwoLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(TwoLayerNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

这里有个小技巧。我在项目中经常用nn.Sequential来简化代码，特别是当网络结构是顺序执行的时候：

class TwoLayerNetSeq(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(TwoLayerNetSeq, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_classes)
        )

    def forward(self, x):
        return self.layers(x)

你想想看，这样是不是清爽多了？但要注意，如果你的网络有分支结构（比如ResNet的残差连接），nn.Sequential就不够用了，你得手动写forward。

3.3 模型保存与加载：state_dict 与 torch.jit.trace

模型训练好了，怎么保存？怎么加载？这里有两种主流方式，用途完全不同。

3.3.1 state_dict：保存参数，灵活加载

state_dict是一个Python字典，里面存了模型每一层的参数（权重和偏置）。这是最常用的保存方式。

保存：

torch.save(model.state_dict(), 'model_weights.pth')

加载：

model = TwoLayerNet(784, 256, 10)  # 先创建模型实例
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # 别忘了切换到评估模式

这里有个关键点，加载state_dict前，你必须先创建一个结构完全相同的模型实例。否则会报错。我曾经在项目中因为改了模型结构但忘了更新加载代码，排查了半天才发现是参数名对不上。

小提示： 保存整个模型（包括结构）可以用 torch.save(model, 'model.pth')，但我不推荐。因为这样会序列化整个模型类，如果代码版本变了，加载时容易出问题。还是用 state_dict 更稳妥。

3.3.2 torch.jit.trace：为部署而生

好了，重点来了。torch.jit.trace是PyTorch模型转TensorRT的关键一步。它会把你的PyTorch模型“追踪”成一个静态的计算图，保存为.pt文件。这个文件不依赖Python代码，可以跨平台部署。

怎么用？很简单：

# 假设 model 已经训练好，且处于 eval 模式
model.eval()
example_input = torch.randn(1, 784)  # 一个示例输入，形状要和模型匹配
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_traced.pt')

加载：

loaded_model = torch.jit.load('model_traced.pt')
output = loaded_model(example_input)

你看，加载时不需要再定义模型结构了，直接加载就行。这就是为部署设计的。

核心要点： torch.jit.trace 要求输入张量的形状是固定的。如果你的模型需要处理动态形状（比如不同长度的序列），trace 可能会出问题。这时候可以考虑用 torch.jit.script，但那是另一个话题了。

3.4 避坑指南：我踩过的那些坑

嗯，这里分享几个我实际项目中遇到的坑，希望能帮你省点时间。

  • 忘记 model.eval()：保存或 trace 前，一定要调用 model.eval()。否则像 Dropout、BatchNorm 这些层的行为会不一致，导致推理结果错误。我曾经有一次 trace 完模型，发现推理结果和训练时不一样，查了半天才发现是没切到 eval 模式。
  • 输入形状不匹配：trace 时给的示例输入，形状必须和实际推理时一致。特别是 batch size 维度，如果你 trace 时用了 batch size=1，那后面推理时也只能用 batch size=1。想支持动态 batch？那得用更高级的技巧。
  • 控制流问题：如果你的 forward 里有 if 语句或循环，trace 只会记录走的那一条路径。比如 if x.shape[0] > 10: ... else: ...，trace 时如果输入 batch size 是 1，那 else 分支就不会被记录。这会导致部署后模型行为异常。

3.5 总结

这一章的内容，说白了就是为后面的TensorRT转换打基础。你只要记住：

  • 用 nn.Module 定义模型，记得调用 super().__init__()
  • 保存参数用 state_dict，部署用 torch.jit.trace
  • trace 前一定要 model.eval()，输入形状要固定

下一章，我们会正式进入TensorRT的世界。到时候你会发现，今天这些基础工作做得越扎实，后面就越顺利。

好，咱们下章见。

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