2. 模型加载与初始化：PyTorch模型加载、权重映射、设备分配（CPU/GPU）
好，咱们正式开始聊推理的第一个实操环节——模型加载与初始化。
很多人觉得这步很简单，不就是 torch.load() 一下嘛？
嗯，我刚开始也这么想。直到有一次在生产环境里，模型死活加载不上，报了一堆 key 不匹配的错误。排查了半天，才发现是 checkpoint 里保存的权重名字跟我定义的模型层名字对不上。
从那以后，我养成了一个习惯：加载模型前，先搞清楚权重的映射关系。
2.1 PyTorch 模型加载的三种方式
PyTorch 官方给了我们三种保存和加载的方式。说白了，就是三种不同的“打包”策略。
| 方式 | 保存内容 | 加载方式 | 适用场景 |
|---|---|---|---|
| 方式一：保存整个模型 | 模型结构 + 权重 | torch.load('model.pth') |
快速实验，不推荐生产 |
| 方式二：保存 state_dict | 仅权重（推荐） | model.load_state_dict(torch.load('weights.pth')) |
生产环境，灵活可控 |
| 方式三：保存 checkpoint | 权重 + 优化器状态 + epoch | 按需加载各部分 | 训练中断恢复 |
我个人习惯用方式二。为什么呢？
因为方式一有个坑：它把模型类也序列化了。如果你换了环境，或者改了模型定义，加载时直接报错。你想想看，推理服务最怕什么？最怕启动时挂掉。
核心原则：推理阶段只加载权重，不加载模型结构。模型结构应该在代码里显式定义。
2.2 权重映射：那些年我踩过的坑
权重映射，说白了就是 checkpoint 里的 key 和 模型里的 key 怎么对上。
举个例子。你训练时用了 DataParallel，模型层名前多了 module. 前缀。推理时你用的是单卡，模型定义里没有这个前缀。直接加载？报错。
我曾经在一个线上事故里排查了整整两个小时，最后发现就是多了一个 module.。
解决方案其实很简单：
# 去掉 'module.' 前缀
state_dict = torch.load('checkpoint.pth')
new_state_dict = {}
for k, v in state_dict.items():
name = k.replace('module.', '') # 去掉前缀
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
小技巧：保存 checkpoint 时，建议统一用 model.module.state_dict() 而不是 model.state_dict()，这样不管用不用 DataParallel，key 的格式都是一致的。
2.3 设备分配：CPU 还是 GPU？
模型加载到哪个设备上，这个决策其实挺关键的。
你可能会说：那肯定 GPU 啊，推理快。嗯，大部分情况是这样。但有一种情况我建议你用 CPU——当你的模型太大，显存放不下时。
我记得有一次做 BERT-large 的推理，单卡 16G 显存根本放不下。怎么办？
两个思路：
- 思路一：用 CPU 推理，慢是慢了点，但能跑。
- 思路二：用模型并行或流水线并行，把模型切分到多张卡上。
对于单机推理，我建议用 torch.device 统一管理设备：
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyModel().to(device)
# 加载权重时也要指定 map_location
state_dict = torch.load('weights.pth', map_location=device)
model.load_state_dict(state_dict)
注意：如果不指定 map_location，PyTorch 默认会把权重加载到保存时的设备上。如果你在 GPU 上保存的权重，在只有 CPU 的机器上加载，直接报错。
2.4 核心流程图：模型加载与初始化全流程
下面这张图，是我自己总结的。每次写推理代码前，我都会对照着走一遍。
2.5 实战：一个完整的加载示例
最后，给你一个我常用的模板。每次写推理代码，我都是从这个模板开始的。
import torch
import torch.nn as nn
class TransformerModel(nn.Module):
def __init__(self, vocab_size, d_model, nhead, num_layers):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.transformer = nn.Transformer(d_model, nhead, num_layers)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = self.transformer(x)
return self.fc(x)
# 1. 定义模型
model = TransformerModel(vocab_size=30000, d_model=512, nhead=8, num_layers=6)
# 2. 加载权重（带映射）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
checkpoint = torch.load('transformer_weights.pth', map_location=device)
# 3. 处理 key 映射（如果有 module. 前缀）
state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
new_state_dict = {}
for k, v in state_dict.items():
name = k.replace('module.', '')
new_state_dict[name] = v
# 4. 加载到模型
model.load_state_dict(new_state_dict, strict=False) # strict=False 允许部分加载
model.to(device)
model.eval() # 切换到推理模式
# 5. 开始推理
with torch.no_grad():
input_ids = torch.randint(0, 30000, (1, 128)).to(device)
output = model(input_ids)
print(output.shape) # torch.Size([1, 128, 30000])
个人建议：加载时用 strict=False，这样即使有些层没加载上，也不会直接崩溃。然后自己打印日志检查哪些 key 没对上。嗯，这个习惯帮我避免了好几次线上事故。
好了，模型加载与初始化就聊到这里。记住一句话：加载模型不是简单的文件读取，而是一个结构对齐、设备分配、状态切换的完整流程。每一步都走稳了，推理才能跑得顺。