Python 现代开发学习笔记

(二十七)恢复模型、优化器与随机状态

# 续写同一文件: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 表示目标总轮数。已完成三轮,再指定六轮,会继续第四、五、六轮;不是额外再跑六轮。

三轮结束写入参数、Adam、随机层与打乱状态,恢复后从第四轮继续,与连续六轮对齐比较相同批次、历史和最终权重

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