双偏振雷达与雨量计融合的短时降水预报：PyTorch数据组织与卷积网络实战

简介基于PyTorch的雷达降水预测模型资源融合雷达反射率、差分反射率、差分相移率及地面雨量计观测等多源气象数据通过自定义数据加载器完成时序数据预处理与归一化并借助卷积神经网络提取空间特征面向气象科研人员、深度学习开发者及有降水预测需求的学生提供了可复现的完整研究样例。压缩包共102个文件含7个Python脚本、3个pth模型权重、5个训练日志、84张预测结果可视化png图以及md说明文档、docx附赠资料和txt配置信息整体约45.73MB目录结构清晰便于按模块学习与二次开发。目前已有128人浏览学习。模型源码、预训练权重、训练日志与可视化结果一应俱全既能帮助快速掌握多源数据预处理、CNN特征提取与模型训练评估流程也可直接作为雷达降水外推或时序预测项目的参考基线。结合项目说明文档与附赠资料读者可快速复现实验并理解模型设计细节适合作为进一步研究或工程落地的起点。1. 短时降水预报为什么卡在数据组织上双偏振量与雨量计怎么进模型我在对接短临预报业务时发现一个很典型的现状真正卡住模型的通常不是网络层数而是数据根本喂不进去。双偏振雷达一次体扫能同时给出反射率因子Z、差分反射率ZDR和差分相移率KDP再加上地面雨量计理论上能把降水粒子尺寸、相态和液态水含量说清楚但它们的量纲、时间分辨率、空间覆盖范围和噪声特性各不一样直接用 np.load 塞进网络训练轻则指标难看重则训练到一半 loss 变成 NaN。基于 PyTorch 把这四类数据组织成可训练的时序样本解决的是从雷达文件到模型输入之间的对齐、归一化和批量加载问题。这篇文章适合正在做雷达降水预测、手里有体扫或格点数据、想把多源气象数据真正用起来的人。我会把数据组织、自定义加载器、卷积网络和评估避坑完整走一遍。2. 训练数据怎么准备Z、ZDR、KDP与雨量计的时序对齐方案2.1 三个偏振量在模型里的角色和各自的预处理重点雷达降水预测常用的是 S 波段或 C 波段双偏振雷达一次体扫能同时拿到三组基础观测量。先不用急着把它们想得很玄学落到模型输入上只需要搞清楚一件事每个量给网络提供什么判别信息以及它有哪些会让训练翻车的噪声。数据量单位典型范围时间分辨率在模型里的角色预处理重点反射率因子 ZdBZ-1060体扫约 6 分钟降水强度的主信号负责“哪里在下”去地物杂波50 dBZ 以上注意衰减区差分反射率 ZDRdB-15同体扫粒子扁率与尺寸区分大雨滴、冰雹与层状云去掉系统偏差异常值 clip差分相移率 KDP°/km-16同体扫液态水含量强降水区比 Z 更抗衰减低信噪比区噪声大必须质控和平滑地面雨量计mm/h 或累计 mm020015 分钟真值标签监督信号来源时间对齐到雷达观测时刻剔除故障站Z 负责降水强度的空间分布ZDR 把“雨滴有多大、是不是冰雹”的信息带进来KDP 在强回波衰减区是比 Z 更可靠的线索——当 Z 被波束衰减吃掉的时候KDP 仍然能认出强降水核心。三者同时进网络等于给模型多一个“即使 Z 被衰减也能从 KDP 恢复强中心”的通道。我一般会先把 KDP 做一次中值滤波再用物理范围 clip因为 KDP 在弱回波区经常跳变成大绝对值这种噪声学到的不是降水特征而是雷达质量的伪影。2.2 时间窗口怎么切、格点怎么对齐、标签怎么定时序样本的组织方式直接决定模型能学到什么。常见做法是把雷达体扫序列切成长度为 T 的输入窗口用过去 T 帧预测未来一段时间的累计降水。体扫间隔约 6 分钟取过去 5 帧就是 30 分钟上下文这对于 3060 分钟的短临预报来说基本够用帧数太少模型看不到降水系统的移动趋势帧数太多则把非降水时段的静止回波也卷进来反而增加假阳性。角色内容说明输入时间窗T-30 到 T 时刻的 5 帧雷达数据每帧包含 Z、ZDR、KDP 三个通道预测目标T30 的未来 30 分钟累计降水由雨量计观测生成单位 mm空间范围160×160 的 1 km 笛卡尔格点覆盖约 160 km 见方的区域标签来源雨量计站点观测插值到格点训练时用站点掩码约束 loss 计算位置极坐标 PPI 数据不能直接进卷积网络需要先重采样到笛卡尔格点。这一步常见的做法是做 1 km 或 0.5 km 分辨率的 CAPPI等高面扫描把同一高度层上的反射率取出来再做坐标变换。标签生成有个容易踩的坑雨量计是稀疏点不可能每个格点都有观测。我一般会先用反距离加权把站点雨量插值到全部格点形成“参考标签”但在计算 loss 时只对有真实站点观测的格点算损失。这样网络仍然能看到空间上下文又不会被插值产生的平滑伪标签过度主导。2.3 数据划分与文件组织按天气过程而不是按样本模型训练时最容易犯的数据泄漏错误是把整段雷达序列随机拆成样本再划分训练集和验证集。雷达帧是强时间相关的同一场降水过程里相邻时刻的样本长得几乎一样如果它们分别落在训练集和验证集验证指标会好得反常上线后立刻被打回原形。我一般会按“天气过程”整体分块先按日期把连续观测分段再整段分配。比如把某几次完整的降水过程全部放在验证集里训练集完全不碰这些时段。文件组织方面我习惯把每个样本预先切好并保存成 npz字段固定为 reflectivity、zdr、kdp、gauge 四个数组文件名带上时间标签。这样自定义数据加载器就只负责读文件、归一化和转 Tensor不用在训练时做重采样这类重活后面排查数据问题也方便。保存格式用 float32不要用 float64否则一个 160×160×5×3 的样本会把显存和内存都撑大好几倍。3. 自定义数据加载器怎么写切窗、归一化与批量输出的PyTorch实现3.1 样本切窗逻辑固定帧数窗口与缺帧处理切窗逻辑看起来简单真正实现时有两个边界要处理。第一个是序列长度不够降水过程初期往往只有零星回波样本可能凑不满 5 帧第二个是体扫缺失雷达故障或数据落盘失败会留下空洞。我的处理方式是优先丢弃最老的一帧来保持窗口对齐到当前时刻如果缺口在窗口中间则用该格点的历史中值填充并在样本里记录一个 valid_mask让网络知道哪些位置是真实观测。切窗时还要注意一个细节Z、ZDR、KDP 三个量来自同一次体扫的同一时刻它们在时间上是天然对齐的但雨量计的观测时刻会稍微滞后于雷达体扫。对齐时我以雷达体扫时刻为基准取该时刻前后各 2 分钟内的雨量计平均值作为当前时段的观测值再往后累计 30 分钟得到标签。这个 2 分钟窗口能有效平滑雨量计的随机抖动又不会把跨时次的降水混进来。3.2 Dataset的完整实现加载、归一化、标签变换一并搞定自定义数据加载器的核心是继承 torch.utils.data.Dataset把“读文件→归一化→构造标签→转 Tensor”全部封装在getitem里。下面是一份可以直接投入使用的实现前提是预处理阶段已经生成了包含五个字段的 npz 文件。import numpy as np import torch from torch.utils.data import Dataset def normalize_z(z): # dBZ 量纲截断到 0~60再平移到 [-1, 1] return (np.clip(z, 0.0, 60.0) - 30.0) / 30.0 def normalize_zdr(zdr): # ZDR 有效范围小clip 后减去中值再缩放 return (np.clip(zdr, -1.0, 5.0) - 2.0) / 3.0 def normalize_kdp(kdp): # KDP 噪声大clip 到物理合理范围后除以最大值 return np.clip(kdp, -1.0, 6.0) / 6.0 class RainSequenceDataset(Dataset): def __init__(self, npz_list, input_frames5, modetrain): self.npz_list npz_list self.input_frames input_frames self.mode mode def __len__(self): return len(self.npz_list) def __getitem__(self, idx): with np.load(self.npz_list[idx]) as d: z d[reflectivity].astype(np.float32) # (T,H,W) zdr d[zdr].astype(np.float32) # (T,H,W) kdp d[kdp].astype(np.float32) # (T,H,W) label_raw d[gauge].astype(np.float32) # (H,W) 原始雨量 station_mask d[station_mask].astype(np.float32) # (H,W) # 取最后 input_frames 帧 z z[-self.input_frames:] zdr zdr[-self.input_frames:] kdp kdp[-self.input_frames:] # 堆叠成 (T, 3, H, W)三个量分别归一化 x np.stack([normalize_z(z), normalize_zdr(zdr), normalize_kdp(kdp)], axis1) x torch.from_numpy(x).float() # 标签做 log1p 压缩抑制大雨量级的极端值 y torch.from_numpy(np.log1p(label_raw)).float().unsqueeze(0) mask torch.from_numpy(station_mask).float().unsqueeze(0) return x, y, mask这份代码把三类抽象都落到了实处其一归一化放在加载器里而不是预处理脚本里换归一化参数不需要重新生成整个数据集其二label 用 log1p 压缩因为雨量计数值从 0 到超过 100 mm/h直接回归会让模型只学“无降水”这一主流类别其三station_mask 随样本一起返回供损失函数只在真实观测点上计算。注意 label 是 30 分钟累计雨量的 log 值预测时反向 expm1 就能得到原始雨量。3.3 归一化参数怎么定从量纲到数值范围的映射三个偏振量的物理单位完全不同如果不归一化直接送入卷积网络第一个卷积层就会让大数值通道比如 dBZ主导梯度。归一化的目标是让每个通道在输入空间大致落在 [-1, 1] 附近这样 BN 层的压力小初始训练也更稳定。通道原始范围归一化公式设计理由Z-1060 dBZ(clip(z,0,60)-30)/30忽略负 dBZ 的杂波区把 30 dBZ 附近的层状云放在零点ZDR-15 dB(clip(zdr,-1,5)-2)/3ZDR 中值约 2 dB大雨滴正偏冰雹接近 0KDP-16 °/kmclip(kdp,-1,6)/6强降水区 KDP 可达 46clip 掉异常尖峰标签0200 mmlog1p(clip(gauge,0,200))log 压缩动态范围降低极大值对 MSE 的统治Z 的 clip 下边界我放在 0 dBZ而不是 -10 dBZ因为负 dBZ 大多是地物杂波或晴空回波对降水预测没有信息量留着只会增加训练样本里的“无意义背景”。KDP 的 clip 上边界 6°/km 是经验值超过这个值的数据基本来自衰减区噪声或融化层异常保留会让模型学到不合理的强中心。这组参数不需要每个区域都调一遍但换雷达型号或换季节时我建议重新扫一遍数据分布再确认 clip 边界。3.4 DataLoader参数与训练提速的配置Dataset 写好后DataLoader 的配置同样影响训练效率。我的常用配置是 num_workers 取 4 到 8pin_memory 打开batch_size 按显存从 8 起试。代码块如下from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size16, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue, ) val_loader DataLoader( val_dataset, batch_size16, shuffleFalse, num_workers4, pin_memoryTrue, )这里 drop_last 在训练时保留 True避免最后一个 batch 样本数不足导致 BN 层统计量抖动。如果训练时 GPU 利用率上不去而 CPU 被打满先看getitem里是不是做了太多实时计算比如 np.load 后还做插值或滤波这些都该在离线预处理完成。我踩过最大的坑是在加载器里放了反距离插值一个 epoch 跑了十几个小时把插值移出后直接降到两小时以内。提示npz 里的数组 shape 必须固定为 (T,H,W)T 不足时在预处理阶段就补齐不要在getitem里动态补帧不然每个 batch 的计算图大小不一致分布式训练时容易出问题。4. 卷积神经网络怎么搭多源时序输入融合与损失函数设计4.1 为什么用卷积编码器而不是直接拉平全连接雷达回波本质上是空间场数据降水系统在相邻格点之间高度相关局部纹理、梯度、尺度信息都要靠卷积才能高效提取。全连接网络把每个像素当成独立特征不仅参数量爆炸还会丢掉空间局部性实际效果远不如卷积。对于时序部分常见选择有三种把多帧通道直接拼起来做 2D 卷积、把帧维也卷进去做 3D 卷积、或者用 ConvLSTM。标题里的“卷积神经网络”落在最稳的方案上输入是 5 帧×3 通道把时间维拼进通道维得到 15 个输入通道用 2D 卷积编码器处理。通道拼接方案在帧数不超过 10 时完全够用而且显存占用可控、实现简单、排查方便。想要更强的时序建模后续可以在编码器顶部替换成 ConvLSTM 或 3D 卷积其余解码器部分不用动。先把通道拼接方案跑通再谈升级是我在这个方向上一贯的推进顺序。4.2 编码器-解码器结构5帧3通道输入到单通道预测模型主体采用编码器-解码器结构输入形状为 (B, T, C, H, W)在 forward 里先合并 T 和 C 两维经过四次下采样提取多尺度特征再用转置卷积上采样并拼接编码器对应层特征最后输出单通道降水预测。代码如下import torch import torch.nn as nn import torch.nn.functional as F class ConvRainNet(nn.Module): def __init__(self, in_frames5, in_channels3, base32): super().__init__() self.in_frames in_frames # 编码器通道数逐级翻倍 enc_channels [in_frames * in_channels, base, base * 2, base * 4, base * 8] self.enc1 self._conv_block(enc_channels[0], enc_channels[1]) self.enc2 self._conv_block(enc_channels[1], enc_channels[2]) self.enc3 self._conv_block(enc_channels[2], enc_channels[3]) self.enc4 self._conv_block(enc_channels[3], enc_channels[4]) # 解码器转置卷积上采样拼接编码器同尺度特征 self.dec3 self._deconv_block(enc_channels[4], enc_channels[3]) self.dec2 self._deconv_block(enc_channels[3] * 2, enc_channels[2]) self.dec1 self._deconv_block(enc_channels[2] * 2, enc_channels[1]) self.head nn.Conv2d(enc_channels[1], 1, kernel_size1) def _conv_block(self, cin, cout): return nn.Sequential( nn.Conv2d(cin, cout, kernel_size3, padding1), nn.BatchNorm2d(cout), nn.ReLU(inplaceTrue), ) def _deconv_block(self, cin, cout): return nn.Sequential( nn.ConvTranspose2d(cin, cout, kernel_size2, stride2), nn.BatchNorm2d(cout), nn.ReLU(inplaceTrue), ) def forward(self, x): # 输入 x: (B, T, C, H, W) B, T, C, H, W x.shape x x.reshape(B, T * C, H, W) # 合并时间维到通道维 e1 self.enc1(x) e1_pool F.max_pool2d(e1, 2) e2 self.enc2(e1_pool) e2_pool F.max_pool2d(e2, 2) e3 self.enc3(e2_pool) e3_pool F.max_pool2d(e3, 2) e4 self.enc4(e3_pool) d3 self.dec3(e4) d3 torch.cat([d3, e3], dim1) # 跳跃连接 d2 self.dec2(d3) d2 torch.cat([d2, e2], dim1) d1 self.dec1(d2) out self.head(d1) out F.interpolate(out, size(H, W), modebilinear, align_cornersFalse) return out.squeeze(1) # 输出 (B, H, W) log 空间的降水预测这段代码里最关键的设计是三次下采样与三次上采样的对称结构加上编码器-解码器之间的跳跃连接。跳跃连接能把低层的高分辨率纹理信息直接送到解码器否则上采样出来的预测图像会模糊成一团强降水中心位置漂移。reshape 操作把时间维并入通道维后卷积核在空间上共享权重同时用不同通道表达不同时刻的状态。最后 interpolate 把输出恢复到原图分辨率方便直接和站点标签做逐格点比较。注意模型输出的是 log1p 空间的数值不是原始雨量。做推理时先 expm1 再取整否则强降水中心会被 log 压缩掩盖。4.3 损失函数加权Huber与强降水权重降水预测的样本极不平衡无降水或弱降水格点占绝大多数直接算 MSE 会让模型学会“秃头预测”所有位置都是零。我的做法是 Huber Loss 加上按雨强分级的权重图对强降水格点给更高的惩罚。def make_weight_map(label_raw, thresholds(0.5, 5.0, 20.0)): w torch.ones_like(label_raw) w[label_raw thresholds[0]] 2.0 # 弱降水 w[label_raw thresholds[1]] 4.0 # 中等降水 w[label_raw thresholds[2]] 8.0 # 强降水 return w def weighted_huber_loss(pred, target, mask, label_raw): diff pred - target # Huber 折中点取 1.0log 空间下 1 对应原始雨量 e-1 ≈ 1.72 mm huber torch.where(diff.abs() 1.0, 0.5 * diff * diff, diff.abs() - 0.5) weight make_weight_map(label_raw) * mask return (huber * weight).sum() / mask.sum()这里 mask 的作用是把 loss 限制在真实雨量计站点的格点上预测图在别处虽然有数值但不算损失。权重设置上强降水格点的 loss 是普通格点的 8 倍这样模型才有动力去拟合极端值而不是输出一片零。折中点在 1.0 是因为标签做了 log1p差值为 1 对应于原始雨量约 1.72 mm 的误差这个尺度下用 L2 会让模型对强降水误差过度敏感用 L1 又不够平滑Huber 是更稳的选择。4.4 超参数一览与显存注意模型层面的超参数不需要太多调优跑通基线后一般只动三个base 通道数、输入帧数、batch_size。参数推荐值调整方向base 通道数32显存不足减到 16效果不够增到 64输入帧数510降水系统移动快时加大帧数batch_size816按显存上限取learning rate1e-3 起Adam训练 20 epoch 后降到 1e-4输入尺寸160×160换分辨率要同步改池化层数显存方面输入是 5 帧×3 通道的 160×160 图大部分情况下 16 GB 显存可以跑通。如果数据分辨率更高建议先用 64×64 跑通验证流程再放大不要一上来就追求全尺寸训练否则定位 bug 的成本会成倍上升。帧数加到 10 时显存占用几乎线性增长必要时把 base 减半来换。5. 训练与评估的避坑清单数据泄漏、KDP噪声与模糊预测5.1 数据侧的三个坑泄漏、NaN与样本失衡坑 1验证指标好得反常上线后立刻失效。原因是样本随机划分导致同一场降水过程的相邻帧同时落在训练集和验证集模型相当于“记住了”验证集样本的邻居。解决办法是改成按日期或按天气过程整体分块验证集完整拿掉某几场降水过程训练集绝不包含这些时间段的数据。我当时的判断依据很简单验证集 loss 降到 0.08 附近但地面实况的强降水中心几乎全漏典型的数据泄漏症状。坑 2训练中途 loss 变成 NaN。根因通常是 KDP 的异常值没处理干净比如融化层里 KDP 出现超过 10 的尖峰log1p 之后反向传播的梯度变得极端。另一个常见来源是 npz 里混入了全零帧归一化后在 0 附近产生无穷大。解决办法是在归一化前统一 clip并对全零帧做检查直接丢弃同时在getitem返回前加一个断言确认 tensor 中没有 NaN 值。坑 3强降水样本太少模型学不到极端情形。雷达降水预测和常规图像回归最大的差别就在这里无降水帧可能占训练集的 70% 以上强降水样本屈指可数。解决办法是给训练集做分层采样按标签里的最大雨强把样本分成弱、中、强三档每档按比例抽取。我在加载器里加了一个简单的档位索引表训练时每个 epoch 先打乱档位顺序再打乱档内样本强降水样本参与训练的次数立刻得到提升。5.2 训练与评估侧的坑CPU瓶颈与模糊预测坑 4GPU 利用率只有 30%CPU 线程全部打满。这是自定义 Dataset 最典型的性能问题原因几乎都是getitem里做了重计算比如实时插值、滤波、数据增强里的随机裁剪。解决办法是把一切能离线算的东西全部离线算完加载器只保留读文件、转类型、切片、归一化四条操作。如果预处理后文件数量过大可以先用 zarr 或多个小文件分片存储但不要在加载器里打散原始数据。坑 5验证集 MSE 不错但画出来的预测图整体模糊强中心被抹平。这是回归任务的通病模型倾向于输出条件均值把所有可能位置都平均到概率场上。MSE 低并不意味着业务可用因为短临预警关心的是“某处未来 30 分钟会不会超过 20 mm/h”而不是逐格点的平均误差。解决办法是评估时改用 TS 评分命中率、空报率和漏报率按 1、5、20 mm/h 三档阈值分开看只报一个 MSE 等于把最重要的问题藏起来了。6. 进阶验证技巧按天气过程回算找到模型的失效边界模型训练完别急着看测试集平均指标我会先做一次“天气过程回算”。选一段完整的强降水过程从过程开始前 30 分钟起每 6 分钟滚动输入一次过去 5 帧输出未来 30 分钟预测再把整段预测序列和雨量计实况做成逐时刻对比动画。这个验证方法能把模型的失效模式暴露得很具体。def roll_eval(model, seq_data, gauge_raw, thresholds(1, 5, 20)): model.eval() all_scores {thr: [] for thr in thresholds} T seq_data.shape[0] step 6 # 体扫间隔单位分钟 for t in range(30, T - 30, step): x torch.from_numpy(seq_data[t - 30:t]).unsqueeze(0) with torch.no_grad(): pred torch.expm1(model(x)).squeeze().numpy() obs gauge_raw[t 30] # 未来 30 分钟实况 for thr in thresholds: pred_bin pred thr obs_bin obs thr hit (pred_bin obs_bin).sum() obs_sum obs_bin.sum() pred_sum pred_bin.sum() ts hit / (pred_sum obs_sum - hit 1e-6) all_scores[thr].append(ts) return {thr: float(np.mean(v)) for thr, v in all_scores.items()}回算之后我习惯把结果按天气阶段分组看层状云降水阶段、对流单体阶段、飑线过境阶段各统计一组 TS 评分。三层分开之后通常会发现模型在层状云阶段表现不错在对流单体阶段漏报率明显偏高。这时候第一反应不是调网络结构而是回头查两件事输入帧数是 5 帧还是太少对流单体从生成到发展成强降水可能只需要 10 分钟KDP 的 clip 边界是不是把强降水的核心信息也一起裁掉了。我在做模拟项目 X 时最早只盯着验证集 loss后来用这种方式回算才发现模型把飑线前缘的窄带回波全漏了把输入帧数从 5 加到 10、放宽 KDP clip 上界之后20 mm/h 档的 TS 提升明显。每一套模型都应该先回答“它在哪类天气上失效”再谈优化这是我的习惯希望帮到你。本文还有配套的精品资源点击获取