ocean-forecast-training

name: ocean-forecast-training description: 海洋时序预测模型训练技能 - 支持多种模型架构的训练、推理与自回归预测（含 OOM 防护 + per-variable 指标 + 续训/切换/恢复工作流） version: 1.2.0 author: Leizheng last_modified: 2026-03-04

海洋时序预测模型训练技能

核心原则

1. 禁止自动决策：模型选择、训练参数、GPU 选择必须由用户确认
2. 错误附带建议：遇到错误时，展示错误信息 + 可能的原因 + 修改建议
3. 错误不自动重试：展示错误分析后询问用户是否调整参数重试
4. 训练完成后主动询问：检测到训练完成时，主动询问是否生成可视化和报告
5. 主动状态感知：训练启动后立即等待，捕获快速完成或早期崩溃

与超分辨率训练的核心区别

可用工具

工作流程

1. 检查 GPU → ocean_forecast_check_gpu
   ↓
2. 确认数据 → 用户提供 dataset_root（预处理输出目录）和 log_dir
   ↓
3. 验证数据 → ocean_forecast_train_start（Stage 1: 自动调用 validate_dataset.py）
   │  展示: dyn_vars, spatial_shape, splits, time_range
   ↓
4. 选择模型 → ocean_forecast_list_models，用户选择
   │  推荐: FNO2d > UNet2d > SwinTransformerV2 > Transformer
   ↓
5. 确认参数 → ocean_forecast_train_start（Stage 3）
   │  时序参数: in_t（默认 7）, out_t（默认 1）, stride（默认 1）
   │  训练参数: epochs, lr, batch_size(默认 4), GPU 选择
   │  OOM 防护: use_amp（默认开启）, gradient_checkpointing（默认开启）
   │  Agent 推荐: in_t 范围 5~14，out_t 范围 1~7，stride 通常 1
   ↓
6. 参数汇总 → 展示所有参数，等待"确认执行"
   ↓
7. 启动训练 → ocean_forecast_train_start（Stage 4, user_confirmed=true）
   │  工具内部等待 training_start 事件（最长 5 分钟）
   │  若返回 status="error"：展示错误 + 建议
   ↓
8. 首次等待 → ocean_forecast_train_status({ action: "wait", process_id, timeout: 120 })
   │  若 completed：主动询问是否生成可视化和报告
   │  若 failed：展示 error_summary + suggestions
   │  若 running（超时）：告知用户训练仍在运行
   ↓
9. 生成可视化 → ocean_forecast_train_visualize（用户确认后）
   │  **禁止在此步骤未成功前调用 ocean_forecast_train_report**
   ↓
10. 生成报告 → ocean_forecast_train_report
    │  Agent 读取报告，补充 <!-- AI_FILL: ... --> 占位符
    ↓
11. 完成 → 展示报告路径和关键结果

预测工作流程（Predict Mode）

predict 模式对测试集执行全量推理，支持自回归 rollout 多步预测。

触发条件

• 用户要求"对测试集做推理/预测/predict"
• 训练完成后需要生成完整预测输出

工作流

1. 确认参数 → dataset_root, log_dir, model_name, ckpt_path（可选）
   ↓
2. 启动推理 → ocean_forecast_train_start({ mode: "predict", dataset_root, log_dir, model_name, ... })
   │  工具内部等待 predict_start 事件（最长 5 分钟）
   ↓
3. 等待完成 → ocean_forecast_train_status({ action: "wait", process_id, timeout: 300 })
   │  若 completed：主动询问是否生成可视化
   │  若 failed：展示错误 + 建议
   ↓
4. 可视化 → ocean_forecast_train_visualize({ log_dir, mode: "predict", dataset_root })
   ↓
5. 完成 → 展示 predictions/ 路径和可视化图表

predict 参数

输出目录

log_dir/
├── predictions/
│   ├── sample_000000_t0_var0_uo.npy
│   ├── sample_000000_t0_var1_vo.npy
│   └── ...
└── plots/
    ├── predict_sample_0_var_uo.png
    ├── predict_overview.png
    └── ...

断点续训工作流

触发条件

• 训练中断（进程崩溃、手动终止、服务器重启）
• 用户要求追加 epoch 继续训练
• 用户要求微调学习率后继续训练

工作流

1. 确认续训条件 → 检查 log_dir 下是否有 best_model.pth
   │  若无：提示用户无可用 checkpoint，需重新训练
   ↓
2. 确认参数 → 使用相同 model_name + ckpt_path 参数
   │  可调整：epochs（追加）、lr（微调）、batch_size
   │  不可变：model_name（模型架构必须一致）、in_t、out_t（时序参数必须一致）
   ↓
3. 参数汇总 → 展示所有参数（标注 ckpt_path），等待"确认执行"
   ↓
4. 启动训练 → ocean_forecast_train_start({ ..., ckpt_path: "{log_dir}/best_model.pth" })
   ↓
5. 后续流程 → 同正常训练（等待 → 可视化 → 报告）

约束

• ckpt_path 必须通过工具参数传入，禁止手动加载
• 模型架构不可变：续训时 model_name 必须与原始训练一致
• 时序参数不可变：续训时 in_t 和 out_t 必须与原始训练一致（影响模型输入输出维度）

切换模型工作流

触发条件

• 用户要求更换模型重新训练
• 当前模型效果不佳，需要尝试其他模型

工作流

1. 终止当前训练 → ocean_forecast_train_status({ action: "kill", process_id })
   ↓
2. 选择新模型 → ocean_forecast_list_models，用户选择
   ↓
3. 新建 log_dir → 新模型必须使用新的 log_dir
   │  禁止：复用旧模型的 log_dir
   ↓
4. 从参数确认步骤开始 → 进入正常训练工作流步骤 5

约束

• 新模型 必须使用新 log_dir，禁止覆盖旧模型的训练输出
• checkpoint 不可跨模型加载（不同模型架构的权重不兼容）

训练失败恢复工作流

触发条件

• 训练返回 process_status="failed"
• ocean_forecast_train_status 返回 error_summary

诊断流程

1. 查看错误摘要 → ocean_forecast_train_status({ action: "status", process_id })
   │  读取 error_summary 和 suggestions
   ↓
2. 查看详细日志 → ocean_forecast_train_status({ action: "logs", process_id, tail: 100 })
   ↓
3. 按决策树分类 → 参考 references/troubleshooting.md 诊断决策树
   │
   ├── 启动崩溃 → 检查 Python 环境 / DDP 配置 / 数据路径
   ├── 训练中途失败 → 检查 OOM / NCCL / NaN / Shape / in_t/out_t
   └── Agent bash 失败 → 立即回到训练工具
   ↓
4. 执行恢复 → 根据诊断结果调整参数，重新启动训练

代码修改规范

• 只改工作空间副本：修改训练代码时只能编辑 {log_dir}/_ocean_forecast_code/ 目录
• 禁止修改原始目录：scripts/ocean-forecast-training/ 目录禁止修改
• 参考 references/troubleshooting.md 的工作空间文件结构和修改权限表

主动状态检查

重要：如果之前启动过训练进程，Agent 在每次收到用户新消息时， 应先调用 ocean_forecast_train_status({ action: "list" }) 检查训练状态。 如果发现训练已完成或失败，优先告知用户训练结果，再处理用户当前请求。

OOM 自动防护机制

训练前自动进行 GPU 显存预估并自动调参，防止训练过程中 OOM 崩溃。

自动防护流程

1. use_amp 默认开启（FFT 模型如 FNO2d、M2NO2d 自动关闭）
2. gradient_checkpointing 默认开启
3. 显存预估 > 85% 时自动降级：
   • 第一步：开启 AMP（如果未开启且模型支持）
   • 第二步：batch_size 减半（直到 1）
   • 最多 5 次尝试
4. 所有手段用尽仍不够 → 报错并建议使用更大显存 GPU

OOM 时的手动建议优先级

1. 启用 use_amp=true（最易操作，效果显著）
2. 减小 batch_size（如 4 → 2 或 1）
3. 启用 gradient_checkpointing=true
4. 减小 in_t（减少输入时间步）
5. 使用多卡训练分摊显存

禁止行为

数据目录要求

需要 ocean-forecast-data-preprocess 预处理后的标准输出目录：

dataset_root/
├── var_names.json          ← 变量名和空间形状
├── time_index.json         ← 时间索引（可选）
├── train/
│   ├── {var_name}/
│   │   ├── 19930101.npy    ← (H, W) float32
│   │   ├── 19930102.npy
│   │   └── ...
│   └── ...
├── valid/
│   └── {var_name}/*.npy
├── test/
│   └── {var_name}/*.npy
└── static/                 (可选)
    └── {var_name}.npy

输出目录结构

log_dir/
├── train.log               ← 训练日志（含 __event__ 标记）
├── config.yaml             ← 训练配置备份
├── best_model.pth          ← 最佳模型权重
├── predictions/            ← predict 模式输出
├── training_report.md      ← 训练报告
└── plots/                  ← 可视化图表
    ├── loss_curve.png
    ├── metrics_curve.png
    ├── lr_curve.png
    ├── per_var_metrics.png
    ├── training_summary.png
    └── predict_*.png

参考文档索引