Python 现代开发学习笔记

(二十六)准备一次训练的状态

# 续写同一文件:training_project/src/readingtorch/engine.py
def train(
    task,
    data,
    output,
    epochs=30,
    seed=7,
    batch_size=24,
    lr=0.03,
    device="cpu",
    resume=None,
):
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    if epochs < 1 or batch_size < 1 or lr <= 0:
        raise ValueError("轮数、批量大小和学习率应大于零")
    if device == "cuda" and not torch.cuda.is_available():
        raise ValueError("当前环境没有可用 CUDA,请选择 --device cpu")
    # 固定本次 CPU 实验的线程与随机状态,恢复时再装回检查点中的状态。
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    torch.manual_seed(seed)
    target_device = torch.device(device)
    bundle = load_bundle(task, data, seed)
    model = build_model(task).to(target_device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    generator = torch.Generator().manual_seed(seed + 1)
    config = {"task": task, "seed": seed, "batch_size": batch_size, "lr": lr}
    history, start, best_score, best_epoch = [], 0, float("inf"), 0
    best_state = snapshot(model)

train 接受任务、数据路径、输出目录和训练设置。它先固定这次 CPU 实验的随机种子,再加载数据、创建模型和 Adam。模型用全局 PyTorch 随机状态初始化;训练批次另用独立的 Generator(seed+1) 安排。

history 保存数值记录,start 表示已经完成多少轮,best_score 初始为无穷大,第一次验证后即可替换。config 把恢复时需要一致的任务、种子、批量大小和学习率放进同一份记录。

模型创建后交给优化器的参数,就是后面 step 修改的那些对象。设备选择也在这里完成,输入批次会在 one_epoch 中移动到相同设备。

从种子和数据划分进入模型初始化与 Adam,训练生成器单独登记,轮数历史与最佳快照共同形成一次训练的起点

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