3、PyTorch模型基础回顾：nn.Module、forward方法、state_dict、模型保存与加载

好，咱们正式开始聊PyTorch模型互转之前，得先把PyTorch这边的基础夯实了。说实话，很多同学一上来就想着怎么把模型转成JAX，结果连PyTorch的state_dict是啥都没搞明白——这就像还没学会走路就想跑马拉松，肯定要摔跟头。

我个人习惯把PyTorch的模型理解成「搭积木」。你想想看，每一块积木（层）怎么拼、怎么连，最后怎么把整个作品存起来，这就是我们今天要聊的核心。

nn.Module：一切模型的基类

在PyTorch里，所有神经网络模型都要继承nn.Module。这不是什么花哨的设计，说白了就是给你一个「框架」，让你往里面填东西。

我刚开始用PyTorch的时候，犯过一个低级错误：直接写了个Python类，里面定义了一堆nn.Linear，但没继承nn.Module。结果模型死活不收敛，查了半天才发现参数根本没注册进去。嗯，这个坑我替你们踩过了。

正确的做法是这样的：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

注意看，super().__init__()这行不能少。它干了件什么事？它帮你在背后注册了所有self.xxx = nn.Linear(...)这样的子模块。没有这行，你的模型就是个空壳子。

核心要点：nn.Module会自动追踪所有赋值给self的nn.Module子对象。但如果你用Python列表或者字典来装它们，它就不认识了。这时候需要用nn.ModuleList或nn.ModuleDict。

forward方法：模型的前向传播逻辑

forward方法定义了数据怎么流过你的模型。你想想看，它就是个「管道」，输入进去，经过各层变换，最后输出结果。

这里有个细节很多人会忽略：forward方法的名字是固定的，不能改成forward_pass或者predict。为什么？因为PyTorch在__call__方法里调用了forward，你改名字它就不认了。

我曾经在项目里见过一个同事，把forward改成了run，然后直接调用model.run(x)——结果梯度全断了。嗯，这就是不读文档的代价。

正确的调用方式有两种：

model = MyModel()
x = torch.randn(32, 784)

# 方式一：直接调用（推荐）
output = model(x)

# 方式二：显式调用forward（不推荐，会跳过一些hook）
output = model.forward(x)

我个人建议永远用方式一。因为__call__除了调用forward，还会触发注册的hooks，这对调试和模型转换都很重要。

state_dict：模型的「身份证」

state_dict是什么？说白了就是一个Python字典，里面存了模型所有可训练参数的名称和对应的张量值。

你可以把它理解成模型的「身份证」——记录了每一层的权重和偏置。我经常用它来做两件事：一是检查模型结构是否正确，二是做模型互转时的数据桥梁。

来看看它长什么样：

model = MyModel()
print(model.state_dict().keys())
# 输出类似：
# odict_keys(['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

# 查看某个参数的值
print(model.state_dict()['fc1.weight'].shape)
# 输出：torch.Size([256, 784])

注意看，这里的key命名规则是层名.参数名。比如fc1.weight，fc1就是你在__init__里定义的属性名，weight是nn.Linear自带的参数名。

小技巧：如果你在模型转换时发现key对不上，八成是层命名不一致。我习惯在定义模型时就统一命名规范，比如用conv1、conv2这样的前缀，方便后续做互转。

模型保存与加载：torch.save和torch.load

模型训练完了，总得存下来吧？PyTorch提供了两种保存方式，我分别说说。

方式一：保存整个模型

# 保存
torch.save(model, 'model.pth')

# 加载
model = torch.load('model.pth')
model.eval()  # 别忘了切换到评估模式

这种方式简单粗暴，但有个大问题：它把模型的结构和参数打包在一起了。如果你换了PyTorch版本，或者改了模型定义，加载时大概率会报错。我个人几乎不用这种方式。

方式二：只保存state_dict（推荐）

# 保存
torch.save(model.state_dict(), 'model_state.pth')

# 加载
model = MyModel()  # 先实例化模型
model.load_state_dict(torch.load('model_state.pth'))
model.eval()

这种方式更灵活。你只需要保证模型结构一致，参数就能正确加载。而且state_dict本身就是个字典，方便你查看、修改，甚至做模型互转。

避坑指南：我曾经在加载模型时忘了调用model.eval()，结果推理结果一直不对。原因是某些层（如Dropout、BatchNorm）在训练和评估模式下行为不同。所以记住：加载完模型后，一定要手动切换到评估模式。

知识体系总览

为了让你更直观地理解这些概念之间的关系，我画了张图：

PyTorch模型基础核心概念 nn.Module（基类） forward() 前向传播 state_dict（参数字典） torch.save / torch.load 关键点 • 继承nn.Module • 必须调用super().__init__() • forward方法名固定 • 用model(x)而非model.forward(x) • state_dict是字典结构 • 推荐只保存state_dict • 加载后调用model.eval() • 注意层命名一致性 • 版本兼容性问题

这张图把整个流程串起来了：从nn.Module开始，定义forward方法，训练后得到state_dict，最后通过torch.save和torch.load来持久化。每一步都环环相扣，缺一不可。

好了，PyTorch模型基础就聊到这儿。这些概念看着简单，但实际用起来门道不少。下一章我们开始正式进入JAX的世界，到时候你会发现，很多思路是相通的，但实现方式完全不同。

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

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