geo-correction 0.1.0

Meteorological grid correction training and inference toolkit

气象 AI 订正训练工程（train/）说明

本文件是基于当前 train/ 目录代码结构整理的新版 README 草案，目标是帮助你快速理解工程组织、主流程、配置方式，以及当前仓库里需要注意的兼容性问题。

1. 项目定位

该工程用于气象格点 AI 订正，覆盖以下典型场景：

• 常规地面要素订正（温度、湿度、风、气压、能见度等）
• 降水订正（TP）
• 高空层（Upar）订正
• 海洋/海浪/海流相关订正
• 在线业务推理（从模式数据生成订正产品）

工程同时包含训练、评估、推理、数据集制作、结果可视化、物理量转换等模块。

2. 工程结构总览

建议优先理解以下目录：

train/
├─ main.py / main_ddp.py                # 常规训练入口（单卡 / DDP）
├─ eval.py / eval_tp.py                 # 离线评估入口（常规/降水）
├─ inference.py                         # 在线推理入口（实时业务链路）
├─ module/
│  ├─ pretrain.py                       # 训练器（Trainer/LRESOTrainer/Ocean...）
│  └─ forecast.py                       # 推理器（WeatherForecast/TPForecast/...）
├─ data/
│  ├─ dataset.py                        # Dataset 与归一化/反归一化逻辑
│  ├─ makeData.py                       # 数据集制作（离线）
│  └─ makeOnlineData.py                 # 在线模式数据读取/拼装
├─ model/                               # 网络结构（AFNO/UNet/HAN/Swin/GraphCast/MOS/...）
├─ loss/                                # 损失函数与指标
├─ tools/
│  ├─ config.py                         # 配置读取 + 环境路径 + 区域定义
│  └─ utils.py                          # 通用工具（索引、归一化读取、插值等）
├─ yaml/                                # 训练/测试 YAML 配置（085/095/...）
├─ json_defs/                           # 数据制作任务定义（JSON）
├─ saved_models/                        # 模型权重（示例）
├─ run/                                 # TensorBoard 日志目录
└─ images/                              # 验证图、误差图等输出（运行后生成）

3. 核心流程（训练 / 评估 / 推理）

3.1 训练流程（常规入口）

主入口是 main.py（常规）和 main_ddp.py（DDP / 多卡）。

训练流程大致为：

1. 读取 YAML 配置
2. 根据 Input_ELE / Output_ELE 从 mean_std_feature.json 计算特征索引
3. 构建 Trainer（陆地或海洋）
4. Trainer 内部创建 Dataset / DataLoader、模型、优化器、损失函数
5. 训练 + 验证 + TensorBoard + 保存最优 checkpoint

代码关系（关键路径）：

• main.py -> module.pretrain.LRESOTrainer / OceanHighWindTrainer
• module/pretrain.py -> data/dataset.py (ECMWFData / OceanData)
• module/pretrain.py -> model/*（AFNO/HAN/Swin/UNet/TransUNet/OceanTransformer/MOS 等）

3.2 评估流程（离线）

• eval.py: 常规、T799、海面大风评估入口
• eval_tp.py: 降水评估入口

两者都会进入 module/forecast.py 的推理器类，逻辑通常包括：

1. 自动查找 saved_models/... 下最佳 checkpoint
2. 加载权重
3. 执行推理
4. 反归一化 / 后处理 / 重塑范围
5. 保存 NetCDF 产品和专题图（默认开启）

3.3 在线推理流程（业务化）

入口为 inference.py，整体流程：

1. 根据业务时间参数（起报时间、时效）读取在线模式数据
2. 通过 data.makeOnlineData 的 RUNNER_REGISTRY 调用对应数据拼装器（ECMWF / T799 / T799Ocean / T1279）
3. 选择 YAML 指定的特征通道
4. 推理预处理（含降水单位处理、归一化）
5. 调用 module.forecast 对应 Forecast 类完成推理与产品落盘

4. 配置体系（YAML）

yaml/ 下有大量按任务拆分的配置文件，例如：

• yaml/085/config_085_HAN_Conver_Train.yaml
• yaml/085/config_085_SWIN_TP_Test.yaml
• yaml/095/config_095_AFNO_Conver_Test.yaml
• yaml/095/config_095_OceanTransUNet_..._Test.yaml

4.1 常见字段

model

• name: 模型名称（如 AFNONet / HAN / SwinTransformerV2Cr / TransUNet / OceanTransformer / MOS-Global）

data

• root_dir: 数据集根目录（通常包含 train/valid/test 子目录）
• statistic_dir: 统计量目录（x_mean.npy/x_std.npy/y_mean.npy/y_std.npy 所在根目录）
• mean_std_feature: 特征与目标索引映射 JSON（Feature/Target）
• AreaType: 区域（如 CHN、095_Area、海洋区域字符串）
• DataType: 任务数据类型（如 Conver / TP / Upar / Ocean）
• TaskType: regression 或 classification
• DZ_Befor_ELE: 订正前要素名列表
• Input_ELE: 模型输入要素名列表
• Output_ELE: 模型输出要素名列表
• grid_resolution: 网格分辨率
• batch_size, num_workers, shuffle
• isnormal: 是否归一化
• scale: 是否降尺度（1 表示不缩放）
• skip_normal_tp: 降水任务可选，跳过 TP 归一化

train_config

• gpu, gpuid, seed
• s_epoch, epoch, lr
• lat_h/lon_w: 输入尺寸
• lat_label_h/lon_label_w: 标签尺寸
• loss_name
• 任务特定字段（如 is_reshape_range, pos_weight, is_logistic 等）

5. 数据集约定（HDF5 / 归一化）

5.1 data/dataset.py 的主要 Dataset

• ECMWFData: 常规任务
• ECMWFDataTP: 降水任务（支持 skip_normal_tp）
• OceanData: 海洋任务（会对 NaN 做 nan_to_num）

5.2 HDF5 常用键（按代码读取逻辑）

输入：

• meteo_ecmwf
• ftime
• dtime

标签（会按顺序自动匹配其一）：

• meteo_era5
• meteo_cldas
• meteo_era5_hycom
• meteo_nbs
• meteo_kd

5.3 归一化与索引

• 训练/评估大量逻辑依赖 mean_std_feature.json 中的 Feature / Target 字段
• tools.utils.get_indices_data() 会根据 YAML 中要素名列表映射到通道索引
• 同时依赖 x_mean.npy/x_std.npy/y_mean.npy/y_std.npy

6. 模型与任务分支（当前代码支持）

model/__init__.py 当前导出的主要模型包括：

• HAN
• SmaAt_UNet
• NestedUNet
• FFNet5
• GraphCast
• AFNONet
• SwinTransformerV2Cr
• TransUNet
• OceanTransformer

训练器与推理器的任务分支主要在：

• module/pretrain.py
• module/forecast.py

其中：

• 常规场：LRESOTrainer / WeatherForecast
• 降水：TPForecast（推理）
• 095 项目：Weather095Forecast
• 海洋大风：OceanHighWindTrainer / OceanHighWindForcast
• Ocean 其他：OceanForcast

7. 输出目录与产物

7.1 训练输出

• saved_models/<数据源>/<区域>/<模型>/<DataType...>/
  • 权重命名示例：AFNONet_0.533_20.pth
• run/<数据源>/<区域>/<模型>/<DataType...>/
  • TensorBoard 日志

7.2 验证可视化输出

• images/vailding/...
  • 单轮可视化图
  • 误差对比图
  • 误差箱线图

7.3 推理产品输出

module/forecast.py 会将推理结果保存为：

• NetCDF (.nc)
• 专题图 (.png)
• 部分场景会追加站点插值 CSV（取决于具体 Forecast 子类）

8. 快速上手（建议版）

8.1 环境依赖（现状）

仓库根部有一个 requirment.sh，目前只写了旧版 PyTorch 安装示例（torch==1.11.0 + cu102），不足以覆盖当前工程依赖。

建议你后续补一个完整 requirements.txt 或 environment.yml。

8.2 训练（常规）

示例（085/095 类配置）：

python main.py --train_config config_095_AFNO_Conver_Train --mode train --area land

海洋训练：

python main.py --train_config config_095_OceanTransUNet_Ocean_hjocean_-15.0N_30.0E-30.0N_100.0E_Train --mode train --area ocean

DDP（海洋大风类训练器）：

torchrun --nproc_per_node=2 main_ddp.py --train_config config_xxx --mode train --area ocean

8.3 评估

常规场：

python eval.py --test_config config_095_AFNO_Conver_Test --mode ecmwf

T799/海洋大风：

python eval.py --test_config config_xxx --mode t799
python eval.py --test_config config_xxx --mode oceanwins

降水：

python eval_tp.py --test_config config_085_SWIN_TP_Test

8.4 在线推理

python inference.py -F 20240101 -H 0 -D 6 -C config_095_AFNO_Conver_Test -R

批量时效模式（代码中使用 -D -1 触发多时效）：

python inference.py -F 20240101 -H 0 -D -1 -C config_xxx -R

9. 已知注意事项（基于当前代码扫描）

以下是我根据当前工程文件确认的事项，建议在正式使用前统一整理：

9.1 启动目录 / 导入路径不一致

部分脚本使用：

• from train.tools.config import load_config

另一些脚本使用：

• from tools.config import load_config

这会导致“从 train/ 目录运行”与“从上级目录把 train 当包运行”之间存在不一致。建议后续统一导入风格（推荐统一为包内相对导入或统一运行入口）。

9.2 tools/config.py 有导入时副作用，强依赖外部环境配置

tools/config.py 在 import 时会立即读取：

• .../env/<env>/application.yml

如果该文件不存在，很多脚本会在 import 阶段直接失败。建议：

• 将环境路径配置与 load_config() 解耦
• 或提供本地默认 application.yml 模板

9.3 部分旧入口脚本与当前 module/pretrain.py 可能已不一致

例如：

• main_tp.py 依赖 TPTrainer, get_criterion
• main_mos.py 依赖 MOSTrainer

但在当前 module/pretrain.py 中未检索到对应定义（看起来像历史重构后的遗留入口）。

同时，main_tp.py / main_mos.py 默认将 YAML 路径拼接为 ./yaml/<config>.yaml，而对应配置多数实际位于 yaml/085/ 或 yaml/095/ 子目录。

9.4 inference.py 存在 095 分支类名未导入风险

inference.py 中使用了 Weather095Forecast 分支，但顶部导入列表里未包含该类，需补充导入后再用。

9.5 data/makeOnlineData.py 的 T799Ocean 分支需再检查

run_t799Ocean_online_data() 中出现了对 np/xr 和 ds_ocean_mask 的使用，但该文件顶部未看到 numpy/xarray 导入，且 ds_ocean_mask 初始化行被注释掉，存在直接运行报错风险。

9.6 module/pretrain.py 的 get_optimizer() 需确认 optim 导入

get_optimizer() 中使用 optim.Adam / optim.AdamW / optim.SGD，建议确认文件顶部是否显式导入 torch.optim as optim（当前扫描结果看起来存在缺失风险）。

9.7 测试脚本可能滞后于当前接口

例如 tests/test_utils.py 中对 get_indices_data() 的调用未传 mean_std_feature，而当前函数实现会检查该路径是否存在。

10. 推荐的后续整理方向（可选）

如果你准备把这个工程对外或给新同事使用，建议优先做这 5 件事：

1. 统一所有入口脚本的导入路径（tools.config vs train.tools.config）
2. 为 tools/config.py 提供本地开发默认配置或延迟加载
3. 清理/修复失效入口（main_tp.py、main_mos.py 等）并统一 YAML 路径策略
4. 补一个可复现环境文件（environment.yml）
5. 增加最小可运行 smoke test（至少覆盖 load_config、Dataset、一个模型前向）