(二十六)准备一次训练的状态
# 续写同一文件: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 中移动到相同设备。