(二十七)恢复模型、优化器与随机状态
# 续写同一文件:training_project/src/readingtorch/engine.py
if resume:
saved = load_checkpoint(resume)
if saved.get("kind") != "training" or saved["config"] != config:
raise ValueError("恢复时请使用原来的任务、种子、批量大小与学习率")
if saved["source_hash"] != bundle.source_hash:
raise ValueError("数据文件已经改变,请另开一次训练")
model.load_state_dict(saved["model"])
optimizer.load_state_dict(saved["optimizer"])
history, start = saved["history"], saved["epoch"]
best_score, best_epoch, best_state = (
saved["best_score"],
saved["best_epoch"],
saved["best_model"],
)
generator.set_state(saved["loader_rng"])
# 接回随机层的下一次抽样位置,模型初始化已在恢复之前完成。
torch.set_rng_state(saved["torch_rng"])
if device == "cuda" and saved["cuda_rng"]:
torch.cuda.set_rng_state_all(saved["cuda_rng"])
if epochs <= start:
raise ValueError("目标总轮数应大于检查点已经完成的轮数")
train_loader = make_loader(bundle, "train", batch_size, generator)
val_loader = make_loader(bundle, "val", batch_size)
loss_fn = nn.MSELoss() if task == "regression" else nn.CrossEntropyLoss()
scale = bundle.preprocess["y_scale"].item() if task == "regression" else 1.0恢复不是重新从头训练。代码读 last.pt,先核对任务设置和数据摘要,再装回当前模型、Adam 的统计、历史、已完成轮数以及已选最佳权重。
Dropout 下一次怎样抽样、下一轮样本怎样打乱,也属于继续执行所需的状态。因此这里恢复 torch_rng 与 loader_rng,使用 CUDA 时还保存和恢复相应的 CUDA 随机状态。
| 状态 | 接回哪一部分过程 |
|---|---|
model |
当前参数和缓冲数据 |
optimizer |
Adam 已积累的更新统计 |
epoch、history |
从下一轮继续,保留已有结果 |
best_model、best_score |
继续比较已经取得的最佳验证结果 |
torch_rng、loader_rng |
接回随机层与批次打乱的位置 |
preprocess、splits、source_hash |
保存输入变换、分组与数据来源 |
epochs 表示目标总轮数。已完成三轮,再指定六轮,会继续第四、五、六轮;不是额外再跑六轮。