5.1 KiB

Raw Permalink Blame History

Checkpoint 保存策略更新说明

更新概述

为了解决硬盘空间限制问题，修改了模型 checkpoint 的保存方式，采用同名覆盖的方式仅保留最后 3 个 epoch的 checkpoint。

主要变更

1. 保存策略

保留的文件

✅ best_model.pt - 最佳模型（当 eval_loss 降低时保存）
✅ latest_checkpoint.pt - 定期保存的检查点（每 save_frequency 步覆盖保存）
✅ epoch_checkpoint_1.pt - 循环保存的最后 3 个 epoch 之一
✅ epoch_checkpoint_2.pt - 循环保存的最后 3 个 epoch 之一
✅ epoch_checkpoint_3.pt - 循环保存的最后 3 个 epoch 之一
✅ epoch_checkpoints.json - epoch checkpoint 元数据文件

移除的文件

❌ final_model.pt - 与 epoch_last 重复，已移除
❌ epoch_*.pt - 不再为每个 epoch 创建独立文件

2. 循环覆盖机制

使用 3 个固定文件名循环覆盖：

Epoch 1  -> epoch_checkpoint_1.pt
Epoch 2  -> epoch_checkpoint_2.pt
Epoch 3  -> epoch_checkpoint_3.pt
Epoch 4  -> epoch_checkpoint_1.pt (覆盖 epoch 1)
Epoch 5  -> epoch_checkpoint_2.pt (覆盖 epoch 2)
...

3. JSON 元数据管理

epoch_checkpoints.json 文件记录：

{
  "checkpoints": [
    {
      "epoch": 8,
      "file": "epoch_checkpoint_2.pt",
      "path": "/path/to/checkpoints/epoch_checkpoint_2.pt",
      "saved_at": "2026-04-11T10:30:00",
      "step": 800
    },
    {
      "epoch": 9,
      "file": "epoch_checkpoint_3.pt",
      "path": "/path/to/checkpoints/epoch_checkpoint_3.pt",
      "saved_at": "2026-04-11T10:35:00",
      "step": 900
    },
    {
      "epoch": 10,
      "file": "epoch_checkpoint_1.pt",
      "path": "/path/to/checkpoints/epoch_checkpoint_1.pt",
      "saved_at": "2026-04-11T10:40:00",
      "step": 1000
    }
  ],
  "next_slot": 1,
  "total_epochs_completed": 10
}

4. 自动恢复功能

训练时会自动从 JSON 文件中读取最新的 checkpoint 并恢复训练：

# 训练时自动恢复
trainer.train(
    resume_from=None,      # 如果指定，优先使用指定的 checkpoint
    reset_training_state=False,
    auto_resume=True       # 自动从最新的 epoch checkpoint 恢复
)

恢复优先级：

如果指定了 resume_from，使用指定的 checkpoint
否则，如果存在 epoch checkpoint 元数据，自动从最新的 epoch checkpoint 恢复
否则，从头开始训练

使用方法

查看保存的 checkpoint

训练结束后，可以通过 JSON 文件查看保存的 checkpoint：

# 查看元数据
cat ./output/checkpoints/epoch_checkpoints.json

# 或者使用 Python
python -c "
import json
with open('./output/checkpoints/epoch_checkpoints.json') as f:
    data = json.load(f)
    print('保存的 epochs:', [cp['epoch'] for cp in data['checkpoints']])
    print('最新的 checkpoint:', max(data['checkpoints'], key=lambda x: x['epoch']))
"

手动加载特定 epoch 的 checkpoint

import torch
from src.model.trainer import Trainer

# 加载元数据
import json
with open('./output/checkpoints/epoch_checkpoints.json') as f:
    data = json.load(f)

# 按 epoch 排序
sorted_checkpoints = sorted(data['checkpoints'], key=lambda x: x['epoch'])
print('可用的 epochs:', [cp['epoch'] for cp in sorted_checkpoints])

# 加载特定 epoch 的 checkpoint
target_epoch = 9
checkpoint_info = next(cp for cp in sorted_checkpoints if cp['epoch'] == target_epoch)
checkpoint_path = checkpoint_info['path']

# 使用 trainer 加载
trainer.load_checkpoint(checkpoint_path)

禁用自动恢复

如果需要从头开始训练，可以禁用自动恢复：

# 在代码中设置
trainer.train(auto_resume=False)

优势

节省磁盘空间 - 只保留最后 3 个 epoch，不会随训练时间增长而占用更多空间
自动管理 - 无需手动删除旧的 checkpoint
顺序清晰 - 通过 JSON 文件可以清楚知道每个 checkpoint 对应的 epoch
自动恢复 - 训练中断后可以从最近的 checkpoint 自动恢复
保留重要 checkpoint - best_model 和 latest_checkpoint 仍然保留

磁盘空间对比

修改前

假设训练 100 个 epoch，每个 checkpoint 100MB:
- epoch_1.pt ~ epoch_100.pt: 100 * 100MB = 10GB
- best_model.pt: 100MB
- final_model.pt: 100MB
- latest_checkpoint.pt: 100MB
总计：~10.3GB

修改后

训练 100 个 epoch，每个 checkpoint 100MB:
- epoch_checkpoint_1~3.pt: 3 * 100MB = 300MB
- best_model.pt: 100MB
- latest_checkpoint.pt: 100MB
- epoch_checkpoints.json: <1KB
总计：~400MB

节省空间：~9.9GB (96% 空间节省)

注意事项

训练中断 - 如果训练在中途被中断，只会保留中断前的最后 3 个 epoch
JSON 文件 - 不要手动修改 epoch_checkpoints.json 文件，否则可能导致恢复失败
兼容性 - 旧的 epoch_*.pt 文件不会自动删除，如果需要可以手动清理

测试

运行测试脚本验证功能：

python test_epoch_checkpoint.py

5.1 KiB Raw Permalink Blame History Unescape Escape