Python 现代开发学习笔记

(二十五)权重快照与检查点文件

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 payload

state_dict() 中的张量会联系到当前模型状态;snapshot 逐项 detach().cpu().clone(),建立独立的 CPU 快照,后面继续训练不会把已选结果一起改掉。

torch.save 写临时文件后,os.replace 用它替换目标文件。读取时显式使用 weights_only=True,恢复张量和普通数据;map_location="cpu" 先把数据装进 CPU,模型随后按需要移动设备。这些机制可查 PyTorch 保存与加载。

当前模型的命名状态经过脱离求导与独立复制形成快照,保存到临时文件后替换,恢复先建结构再按名称装入

笔记正文与全部题解 · Ctrl / ⌘ K