ocean-sr-training

name: ocean-SR-training description: 海洋超分辨率模型训练技能 - 支持多种模型架构的训练、测试、推理（含陆地掩码处理 + OOM 自动防护 + 错误实时反馈 + predict 全图预测 + 续训/切换/恢复工作流） version: 4.4.0 author: Leizheng contributors: kongzhiquan last_modified: 2026-03-04

海洋超分辨率模型训练技能

核心原则

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

可用工具

工作流程

1. 确认数据 → 用户提供预处理数据目录和输出目录
   ↓
2. 选择模型 → ocean_sr_list_models，用户选择
   ↓
3. 确认参数 → epochs, lr, batch_size(默认4), GPU 选择
   │  → OOM 防护参数: use_amp（非 FFT 默认开启 / FFT 默认关闭）, gradient_checkpointing（默认开启）, patch_size
   ↓
4. 参数汇总 → 展示所有参数，等待"确认执行"
   ↓
5. 启动训练 → ocean_sr_train（含事件驱动启动监控）
   │  工具内部等待 training_start 事件（最长 5 分钟）
   │  若返回 status="error"：展示错误 + 建议，询问用户是否调整参数重试
   ↓
6. 首次等待 → ocean_sr_train_status({ action: "wait", process_id, timeout: 120 })
   │  等 2 分钟：
   │  若 process_status="completed"：主动询问是否生成可视化和报告
   │  若 process_status="failed"：展示 error_summary + suggestions
   │    → 询问用户："训练失败原因是 XXX，建议 YYY，是否调整参数重试？"
   │  若 process_status="running"（超时）：
   │    → 告知用户：训练仍在运行中（已完成 N 个 epoch），稍后再询问进度
   ↓
7. 后续询问 → 用户询问"训练怎么样了？"时
   │  调用 ocean_sr_train_status({ action: "wait", process_id, timeout: 120 })
   │  同样等 2 分钟，按上述逻辑处理
   ↓
8. 生成可视化 → ocean_sr_train_visualize（用户确认后）
   │  检查返回 status="success" 且 plots/ 目录下文件已生成
   │  若失败：展示错误，询问用户是否重试
   │  **禁止在此步骤未成功前调用 ocean_sr_train_report**
   ↓
9. 生成报告 → ocean_sr_train_report（仅在步骤 8 成功后执行）
   │  → Agent 读取报告，补充 <!-- AI_FILL: ... --> 占位符
   ↓
10. 完成 → 展示报告路径和关键结果

预测工作流程（Predict Mode）

predict 模式对测试集执行全图 SR 推理，输出物理值空间的 NPY 文件。 跳过训练工作流的 4 阶段确认（OOM/shape/FFT 检查），直接启动。

触发条件

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

工作流

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

predict 参数

输出目录

log_dir/
├── predictions/
│   ├── {filename}_sr.npy    ← 全图 SR 输出 [H, W, C]
│   └── ...
├── test_samples.npz          ← LR/SR/HR 对比数据（前 4 条）
└── plots/
    ├── predict_comparison_00.png ← 可视化对比图
    └── ...

注意事项

• predict 不需要训练参数确认（不走 4 阶段状态机）
• predict 始终使用单卡推理
• 如果 log_dir 中没有 best_model.pth，必须显式传入 ckpt_path

断点续训工作流

触发条件

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

工作流

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

约束

• ckpt_path 必须通过工具参数传入，禁止手动加载
• 模型架构不可变：续训时 model_name 必须与原始训练一致
• 若原始训练使用了 patch_size，续训时应保持相同设置

切换模型工作流

触发条件

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

工作流

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

约束

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

训练失败恢复工作流

触发条件

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

诊断流程

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

代码修改规范

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

主动状态检查

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

OOM 自动防护机制

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

自动防护流程

1. use_amp 按模型默认：非 FFT 默认开启；FFT/频域模型默认关闭（可手动 override）
2. 显存预估 > 85% 时自动降级：
   • 第一步：开启 AMP（如果未开启）
   • 第二步：batch_size 减半（直到 1）
   • 最多 5 次尝试
3. 所有手段用尽仍不够 → 报错并建议使用更大显存 GPU 或设置 patch_size

OOM 时的手动建议优先级

1. 启用 use_amp=true（最易操作，效果显著）
2. 减小 batch_size（如 4 → 2 或 1）
3. 启用 gradient_checkpointing=true（有计算代价；默认已开启，可显式确认）
4. 设置 patch_size（如 64 或 128）
5. 使用多卡训练分摊显存

禁止行为

数据目录要求

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

dataset_root/
├── train/
│   ├── hr/{var}/*.npy
│   └── lr/{var}/*.npy
├── valid/
│   ├── hr/{var}/*.npy
│   └── lr/{var}/*.npy
├── test/
│   ├── hr/{var}/*.npy
│   └── lr/{var}/*.npy
└── static_variables/     (可选)

输出目录结构

训练输出直接保存到 log_dir 指定的目录：

log_dir/                       ← 训练输出目录（由配置指定）
├── train-xxx.log              ← 进程日志
├── train-xxx.error.log        ← 错误日志
├── config.yaml                ← 训练配置
├── train.log                  ← 训练日志
├── best_model.pth             ← 最佳模型权重
├── test_samples.npz           ← 测试/推理样本数据
├── training_report.md         ← 训练报告
├── predictions/               ← predict 模式输出
│   ├── {filename}_sr.npy
│   └── ...
└── plots/                     ← 可视化图表
    ├── loss_curve.png
    ├── metrics_curve.png
    ├── lr_curve.png
    ├── metrics_comparison.png
    ├── training_summary.png
    ├── sample_comparison.png
    └── predict_comparison_XX.png ← predict 模式

参考文档索引