简介这是一份基于多分支融合网络的卫星影像降水预报方案源自2022年NeurIPS Weather4cast竞赛的获奖作品适合深度学习、遥感或气象方向的学生用于课程设计、毕业设计及竞赛复现。压缩包共32个文件约272.8MB以Python源码为主体20个py文件涵盖模型训练、评估与提交脚本另含YAML/CSV配置与数据划分、Jupyter可视化案例及预训练模型权重便于端到端跑通流程。内含运行说明文档、权重文件与示例数据可通过train.py、sat2rad-train.py等脚本训练多分支融合模型并结合可视化notebook分析预测效果从数据加载、模型构建到结果提交的完整链路均有对应代码支撑。目前已有40人学习下载适合具备一定Python和深度学习基础、希望完整复现气象预报任务的开发者借鉴。1. 多分支融合网络卫星影像降水多小时预测这套获奖代码到底能跑多远做气象雷达降水外推的同行应该都有体会短临降水预报一直被当成一个黑匣子问题输入是一段卫星云图序列输出是未来几小时的降水场难在时间跨度一拉长图像细节就糊成一片。这套代码是 2022 年 NeurIPS Weather4cast 竞赛的获奖方案核心思路是拿多分支融合网络同时吃卫星多通道数据把不同时间尺度的特征分开提取再融合最终输出未来 1 到 6 小时的降水预测。我拆完整个包之后的感觉是它不是一个玩具 demo而是能直接换上自己的数据去训练、评估、生成提交文件的完整工程。适合做深度学习课程设计、毕业设计也适合想入口气象 AI 方向的研究生照着复现。Python 基础过得去、能跑通 PyTorch 的人按运行说明一步步来基本能做到端到端跑通。2. 架构与选型多分支融合网络怎么把卫星云图变成降水场2.1 两个训练入口sat-phydnet-train.py 和 sat2rad-train.py 分别做的是什么事打开压缩包第一眼最容易被注意到的就是两个训练脚本sat-phydnet-train.py和sat2rad-train.py。刚开始我以为是重复代码逐行比过之后才发现这两个入口对应的是两种不同的建模思路。sat2rad-train.py走的是纯粹的图像到图像翻译路线输入连续几帧卫星云图输出对应时刻的雷达降水估计。这个思路本质上把降水外推当成一个有监督的序列预测任务输入输出都是稠密图像网络只需要学会从云图外观映射到降水强度。sat-phydnet-train.py则不一样它引入了物理约束分支。这个设计在比赛里比较讨巧因为纯数据驱动模型在降水强度极端值上往往会回归到均值附近而物理引导可以让中间特征更贴近大气运动的基本规律。这里说的物理约束不是真的去跑数值模式而是在网络结构里加了一个可微分的运动场估计模块让模型显式学习云团运动的位移场相当于给网络加了一个结构先验。# 两个训练脚本共同的参数入口以 sat2rad-train.py 为例 parser.add_argument(--model, typestr, defaultmultibranch_fusion, help模型结构选择本仓库默认多分支融合网络) parser.add_argument(--loss, typestr, defaultmse, help损失函数可选 mse / mae / combined) parser.add_argument(--eval_interval, typeint, default1, help每多少个 epoch 跑一次验证集评估) parser.add_argument(--ckpt_dir, typestr, defaultcheckpoints, help权重保存目录注意提前建好目录否则保存时报错)这段代码是入口参数的标配写法。model参数决定了加载哪个网络结构文件换模型时不需要改训练脚本主体loss参数值得多说一句比赛里大多数队伍用的是 MSE 或者 MSE 和 MAE 的组合如果后期发现预测的降水场过于平滑可以考虑换成结构相似性相关的损失这属于进阶调整下文会再提到。2.2 模型结构与主干models.py、base_module.py、backbones 里的设计选择models.py是模型定义的主文件base_module.py提供基础模块backbones目录里放的是主干网络。整体结构走的是编码器-融合模块-解码器路线。编码器把多帧卫星影像编码成多尺度特征融合模块对不同时间尺度的特征做注意力融合解码器逐步上采样恢复到原始分辨率。多分支这个“多”体现在两个维度一是输入多帧二是特征多尺度。帧数维度上模型每次吃进去的是序列中的若干连续帧尺度维度上编码器输出的是多个分辨率的特征图浅层特征保空间细节深层特征保语义信息。在训练时相邻帧之间还会穿插做时间维度的下采样让不同分支接收到不同的时间跨度。class MultiBranchFusion(nn.Module): def __init__(self, backboneresnet34, in_channels12, hidden_dim256): super().__init__() # 每个时间分支共享主干权重这是参数量控制的关键 self.backbone build_backbone(backbone, in_channelsin_channels) self.fusion CrossBranchAttention(hidden_dimhidden_dim) self.decoder nn.Sequential( nn.ConvTranspose2d(hidden_dim, 128, 4, stride2, padding1), nn.ReLU(inplaceTrue), nn.ConvTranspose2d(128, 64, 4, stride2, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(64, 1, 3, padding1) # 输出单通道降水强度场 ) def forward(self, x): # x 的形状: [batch, frames, channels, H, W] frames x.unbind(dim1) feats [self.backbone(f) for f in frames] # 多分支融合输出长度的单位是小时所以把帧数先统一到固定维度 feat self.fusion(feats) return self.decoder(feat)这里最值得注意的设计是主干权重共享。如果不共享12 帧输入就得装 12 份 backbone 参数显存直接爆炸。共享权重之后不同帧用同一套参数提取特征再由融合模块做交叉注意力这是时空序列建模里常见的参数效率做法。backbones目录下默认给了 ResNet 系列的配置也预留了换 Swin Transformer 的接口。比赛阶段很多队伍试过纯 Transformer 结构效果不一定比 ResNet 好原因在于气象数据样本量不够大Transformer 在中小数据集上如果没有充分的预训练反而不容易收敛。这里不是说你不能换而是换之前要想清楚数据量撑不撑得住。2.3 推理脚本 predict-submission.py从保存权重到生成降水预测文件训练完成之后真正要产出的是一个能提交的预测结果文件。predict-submission.py干的就是这件事。它加载训练好的权重读入指定时间段的卫星影像把预测的多个小时结果写成一个数组文件。# predict-submission.py 的核心流程 model build_model(cfg) model.load_state_dict(torch.load(ckpt_path, map_locationcpu)) model model.to(device).eval() with torch.no_grad(): for batch in test_loader: # input_frames 的形状是 [batch, frames, channels, H, W] preds model(batch[input_frames]) save_preds(preds, batch[sample_ids], output_dir)这段代码的细节在save_preds里。项目要求提交文件的 ID 必须和测试集样本一一对应顺序不能乱。很多人训练没问题最后提交时得分极低一查发现是样本 ID 排序和预测结果没对齐。这个脚本的好处是它把样本 ID 和预测值一起写出去了比纯输出张量要稳得多。推理阶段还有一个常见坑训练时用了多尺度融合或者随机裁剪推理时要保证输入尺寸和训练一致。我在复现时发现这个脚本默认不做测试时增强如果你想要更平滑的预测结果可以手动把输入做水平翻转推理再取平均代码里没有内置要自己加。3. 数据与预处理2019 年时间戳划分、缓存样本与 dataloader 的数据流3.1 timestamps_and_splits 文件训练验证集划分的逻辑压缩包里的timestamps_and_splits_2019.csv和timestamps_and_splits_stage2.csv是比赛官方给出的时间戳划分文件。为什么要单独给时间戳而不是直接给文件列表因为气象数据是强时序的划分训练集和验证集时不能随机打乱。如果用普通机器学习那种随机划分模型会看到未来时刻的数据验证集结果会虚高放到真实预测场景立刻露馅。timestamp,split 201901010000,train 201901010030,train 201901020000,validation 201901020030,test这种文件格式很简单但信息量不小。split列标记了每个时刻是训练、验证还是测试。比赛中标准做法是按连续时间段划分比如前 70% 时间训练中间 15% 验证最后 15% 测试。温带气旋的季节性很强这么做是为了模拟真实部署时“用过去预测未来”的场景。加载这个文件时要注意时区问题。文件名里的时间戳是 UTC 还是本地时间直接决定了帧序列对齐是否正确。项目里没有明确说但按 Weather4cast 惯例时间戳是 UTC。如果你打算换成自己的数据建议统一转成 UTC 时间戳并写进配置避免后面手工对齐时出现 8 小时偏差。3.2 w4c_dataloader.py卫星影像如何切块、归一化、按时间组装w4c_dataloader.py是数据加载的核心文件。它做的事情可以拆成四步定位时间戳、读取对应帧的卫星通道、做归一化、按序列窗口组装成张量。class W4CDataset(Dataset): def __init__(self, data_root, split_csv, frames_in12, frames_pred6): self.samples load_timestamps(split_csv) # 逐样本滑动窗口每 12 帧作为输入预测后续 6 帧降水 self.valid_indices self._build_valid_indices() def __getitem__(self, idx): input_frames self._read_frames(self.valid_indices[idx]) target_frames self._read_frames(self.valid_indices[idx] frames_in) # 归一化到 [0, 1]降水数据要保留极值信息不能简单 min-max input_frames normalize_satellite(input_frames) target_frames normalize_rain(target_frames) return {input_frames: input_frames, target_frames: target_frames}这里两个细节值得在意。第一是滑动窗口的步长代码默认每次滑动一个时间步这样训练样本数能覆盖整个时间段数据增强也天然带了时序重叠第二是归一化函数是分开写的卫星影像归一化和降水场归一化用的不是同一套逻辑这是因为卫星通道是反射率数值降水场则带有明显的偏态分布降水强度大的像素虽然少但信息权重极高用统一的 min-max 会把小雨和暴雨压到同一个尺度里模型很难区分。我复现时擅自加了一步对降水目标做了 log1p 变换让偏态分布更接近正态。这一步对后续 MSE 损失的训练有明显帮助但要注意的是评估时要把预测值反变换回去再算指标否则数值会偏小。3.3 cached_sample_ids 与可视化先看数据再调模型cached_sample_ids目录存的是已经处理过的样本 ID 列表作用是避免重复扫描整个时间戳文件。气象数据量很大每次启动训练都重新扫描一遍文件系统会很慢这个缓存机制会把已经验证可用的样本 ID 序列化到磁盘第二次运行直接读取。visualizations.ipynb和sat2rad-visualize.ipynb是数据可视化脚本我的习惯是调模型之前先跑一遍可视化把输入云图和目标降水场并排画出来确认两件事——降水区域和云图高反射率区域的对应关系以及数据 loader 读出来的帧顺序是否正确。# 可视化数据样本的核心逻辑 import matplotlib.pyplot as plt fig, axes plt.subplots(2, 6, figsize(18, 6)) for i in range(6): axes[0, i].imshow(input_frames[i].transpose(1, 2, 0)[:, :, 0]) axes[1, i].imshow(target_frames[i].squeeze(), cmaprain)如果两张图在空间位置上对不上大概率是数据对齐的问题而不是模型的问题这时候去调网络结构纯属浪费力气。这个 notebook 里还画了预测结果和真值的对比对调整超参数很有参考意义。4. 训练与评估从 environment.yml 到 evaluate.py 的完整复现路线4.1 environment.yml 与运行说明先把环境装到能跑为止environment.yml用的是 Conda 环境描述格式定义 Python 版本、PyTorch 版本和依赖库版本。直接 conda env create -f environment.yml 就能把环境建好这是整个复现流程里最不该出问题的一环但实际上很多人在这一步就卡住了。name: weather4cast channels: - pytorch - conda-forge dependencies: - python3.9 - pytorch1.12 - torchvision - numpy - pandas - xarray - netcdf4 - scikit-learn - matplotlib - jupyterlab如果机器是纯 CPU 环境跑推理还能应付训练就别指望了。这个模型的参数量在千万级多帧输入做前向传播和梯度回传至少需要 6GB 显存建议直接用 GPU 服务器。环境装完后先跑一个最小样本的前向测试确认 PyTorch 能正常调用 GPU再开始完整训练省得训练到一半发现显存不够。4.2 train.py 与 sat2rad-train.py参数怎么设训练日志看什么跑训练之前先看运行说明.md这个文件把数据集路径、预训练权重路径和关键超参都写得比较清楚。但它默认你已经有比赛数据如果你是从零开始需要先去官网申请数据下载要不然 loader 阶段就会因为文件缺失直接报错。python train.py \ --data_dir /path/to/data \ --split_csv timestamps_and_splits_2019.csv \ --epochs 100 \ --batch_size 8 \ --lr 0.001 \ --ckpt_dir checkpoints我复现时跑的是 sat2rad-train.pybatch size 从 8 调到 16 之后显存不够用退回 8 就稳了。学习率 0.001 配合 Adam 是这类图像生成任务的常配但如果发现损失曲线在前 10 个 epoch 不下降别急着调学习率先去看数据归一化和 loader 输出是否正确八成是数据问题。训练日志里除了 loss 值还会周期性打印验证集指标。观察的重点是“loss 还在降但验证指标不动了”这种时候多半是过拟合解决办法是提前停止或者加数据增强。比赛代码里一般不会给你写 early stopping需要自己在训练脚本外层套一个回调。4.3 evaluate.py量化降水预测的流程与指标evaluate.py负责计算预测结果和真值的量化指标。Weather4cast 比赛使用的核心指标是对数误差相关的综合打分不是简单的 MSE。评估脚本会逐小时计算预测误差然后按不同预报时效分组汇总。# evaluate.py 中计算指标的简化逻辑 def evaluate_predictions(preds, targets, hours[1, 3, 6]): results {} for h in hours: # 取第 h 小时的预测和真值计算均方根误差 rmse compute_rmse(preds[h], targets[h]) results[f{h}h_rmse] rmse return results一个容易忽略的点降雨是稀疏事件且最大值常常集中在极小区域。RMSE 对大雨中心的误差极其敏感模型稍微把强降雨中心位置偏移几个像素RMSE 就会飙升。所以在看评估结果时不要只看总指标按小时拆开看会发现一个很典型的曲线1 小时预测误差低6 小时预测误差明显变高——这是外推模型的天花板不是 bug。5. 避坑与常见问题预处理、显存、权重加载与平台依赖的五个真实坑5.1 现象一跑就崩报错 KeyError 或者维度对不上第一次跑训练脚本报错信息是KeyError: xxx或者RuntimeError: expected 4D input。原因几乎都是数据路径配置错了loader 没读到正确的帧数。气象数据目录结构比较啰嗦年份、月份、变量各占一层目录脚本里的相对路径和实际路径往往对不上。解决先打印 dataset 返回单个样本的形状确认拿到的是 [12, channels, H, W] 而不是 [channels, H, W]。这个检查只需要加一行 print但如果跳过去后面所有维度报错都会让你绕远路。5.2 现象显存不够batch 设小后 Loss 反而震荡显存不够时把 batch size 从 16 改成 4训练确实能跑但 loss 曲线震荡幅度变大收敛速度明显变慢。批大小直接影响 BatchNorm 统计量的稳定性batch 越小统计噪声越大。解决先检查输入分辨率能否等比缩小。如果原本输入是 256×256可以试着降到 192×192显存占用按平方下降精度损失通常在可接受范围内。如果必须保持分辨率则改用梯度累积每 4 个 batch 更新一次权重模拟 batch size16 的效果。5.3 现象预训练权重加载时报错size mismatchweights.zip解压后加载权重报错信息是 size mismatch通常是主干网络名称对不上。比赛的权重是基于 ResNet 预训练初始化再微调的而本地模型默认的 backbone 名称如果是resnet34权重文件里对应的键名可能是resnet34_base。解决加载权重时忽略不匹配的键然后用缺少的键做随机初始化。对于主干网络之外的多分支融合模块本来就是随机初始化的权重文件里有没有都不影响。5.4 现象数据路径与本机不一致训练脚本半天找不到文件比赛数据目录结构和本机不同比如脚本里写的是/data/train你的数据放在/home/user/dataset下。直接改data_dir参数有时不够因为 loader 内部可能还拼接了硬编码的相对路径。解决打开w4c_dataloader.py检查所有文件路径拼接逻辑。我一般会把所有路径统一改到配置文件的 data_root 下而不是散落在各个脚本里。项目里有些脚本有路径硬编码尤其是可视化 notebook 中几乎必改。5.5 现象在 .ipynb 里跑训练经常 kernel 断开sat2rad-visualize.ipynb这类 notebook 用来跑推理演示没问题但如果直接在里面跑完整训练显存占用持续不释放kernel 会被操作系统杀掉。notebook 环境对长时间运行的训练任务本来就不友好。解决训练脚本用命令行跑notebook 只做数据可视化和结果分析。如果一定要在 notebook 里跑验证每个 cell 结束之后主动释放显存torch.cuda.empty_cache()在这时候能救急但不解决根本问题。6. 更进一步用 visualizations.ipynb 定位降雨漏报再改损失函数提升高分值区间训练完第一版模型之后最值得花时间的不是盲目调学习率而是仔细看可视化结果。打开visualizations.ipynb把验证集里预测和真值并排画出来会注意到两类典型错误一类是强降雨中心位置偏移另一类是弱降雨区域被直接抹掉。弱降雨漏报在视觉上似乎影响不大但降水评分指标对不同强度阈值有加权小到中雨的漏报累计起来对指标影响很显著。这时候有两个改造思路。第一个思路是加一个局部损失权重在降水强度高的像素上提高 MSE 的权重把模型注意力压到强降雨区域。第二个思路是改损失函数为 Focal Loss 风格。Focal Loss 原本用于目标检测里的类别不平衡问题但它的思想对稀疏强降雨预测同样适用让模型专注在难以预测的大值像素上而不是被大量的无降雨背景主导梯度。class FocalRainLoss(nn.Module): def __init__(self, alpha0.5, gamma2.0): super().__init__() self.alpha alpha self.gamma gamma def forward(self, pred, target): # 对降水强度做 log1p 变换压缩动态范围 pred torch.log1p(pred) target torch.log1p(target) diff torch.abs(pred - target) # 大误差像素给予更高权重gamma 控制聚焦程度 weight (torch.ones_like(diff) self.alpha * diff) ** self.gamma return torch.mean(weight * diff ** 2)这个损失把误差动态范围拉开后你会发现模型对强降雨中心的捕捉明显变好代价是中雨区的误差略有上升。整体指标不一定涨但比赛评分往往更看重极端降雨的捕捉能力这部分收益是值得的。换损失之后要把学习率调低一个量级Focal 类损失在数值上比普通 MSE 更大不调学习率容易震荡。从那以后我每次训练完第一版模型都会先让可视化跑一遍盯着强降雨中心看再决定要不要动损失函数。这个习惯帮我省掉了至少三轮无效调参也让我确认了一个道理气象影像模型的瓶颈往往不在网络结构而在损失函数和数据处理对稀疏强信号的态度。希望帮到你。本文还有配套的精品资源点击获取
