(二十五)权重快照与检查点文件
state_dict 是名称到参数、缓冲数据的映射。先用一个最小往返核对:
# 项目文件:demos/state_roundtrip.py
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
torch.manual_seed(7)
model = nn.Linear(2, 1)
x = torch.tensor([[1.0, 2.0]])
before = model(x).detach()
with TemporaryDirectory() as folder:
path = Path(folder) / "weights.pt"
torch.save(model.state_dict(), path)
restored = nn.Linear(2, 1) # 先创建相同结构,再装回命名权重。
restored.load_state_dict(torch.load(path, map_location="cpu", weights_only=True))
restored.eval()
with torch.inference_mode():
print("重载预测相同", torch.equal(before, restored(x)))恢复时先建立相同模型结构,再 load_state_dict 按名称装入。运行后预测逐项相同。权重文件表达模型的数值状态,模型类仍由项目源码提供。
完整项目还需要保存最佳权重和继续训练的状态:
# 项目文件:training_project/src/readingtorch/checkpoint.py
import os
from pathlib import Path
import torch
def snapshot(model):
# clone 创建独立快照,后续训练不会改动已经选中的权重。
return {
name: value.detach().cpu().clone() for name, value in model.state_dict().items()
}
def save_checkpoint(path, payload):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
torch.save(payload, temporary)
os.replace(temporary, path)
def load_checkpoint(path):
# 文件只恢复张量和普通数据,先装入 CPU,再选择运行设备。
payload = torch.load(path, map_location="cpu", weights_only=True)
if payload.get("format") != 1:
raise ValueError("无法识别这份检查点格式")
return payloadstate_dict() 中的张量会联系到当前模型状态;snapshot 逐项 detach().cpu().clone(),建立独立的 CPU 快照,后面继续训练不会把已选结果一起改掉。
torch.save 写临时文件后,os.replace 用它替换目标文件。读取时显式使用 weights_only=True,恢复张量和普通数据;map_location="cpu" 先把数据装进 CPU,模型随后按需要移动设备。这些机制可查 PyTorch 保存与加载。