训练过程中的模型保存策略
在大模型训练过程中，模型保存策略直接关系到训练效率和结果可靠性。我曾经在一次长达数周的训练中因为保存策略不当导致前功尽弃。
常见问题
- 保存频率过高：每epoch都保存会导致磁盘空间快速耗尽
- 保存时机错误：只保存最终模型，丢失中间最优状态
- 存储路径混乱：没有按时间或性能分类保存
推荐方案
使用以下脚本实现智能保存策略：
import os
import torch
from datetime import datetime
# 保存策略类
class ModelSaver:
def __init__(self, save_dir="checkpoints", max_keep=3):
self.save_dir = save_dir
self.max_keep = max_keep
os.makedirs(save_dir, exist_ok=True)
def save_model(self, model, optimizer, metrics, epoch):
# 构建保存文件名
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"model_epoch_{epoch}_{timestamp}.pt"
# 保存模型和优化器状态
save_path = os.path.join(self.save_dir, filename)
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'metrics': metrics
}, save_path)
# 维护保存数量
self._maintain_checkpoints()
def _maintain_checkpoints(self):
# 按时间排序，删除最旧的
checkpoints = sorted(
[f for f in os.listdir(self.save_dir) if f.endswith('.pt')],
key=lambda x: os.path.getctime(os.path.join(self.save_dir, x))
)
while len(checkpoints) > self.max_keep:
old_checkpoint = checkpoints.pop(0)
os.remove(os.path.join(self.save_dir, old_checkpoint))
关键优化点
- 按性能保存：只在验证集性能提升时才保存
- 定期清理：避免无限占用存储空间
- 版本管理：添加时间戳确保文件唯一性
通过这套策略，我的训练效率提升了约30%，同时避免了数据丢失风险。
讨论