最近在跑一个很有意思的项目,核心方向是“任意尺度大气降尺度”,模型采用混合隐式神经网络,并且训练过程基于固定尺度。如果你关注过气象领域的统计降尺度、超分辨率重建,或者接触过隐式神经表示(INR),应该能嗅到这几个词放在一起有多大的信息量。这篇文章我会把这个项目的思路、方案选型、实操过程和踩坑记录完整拆开讲。它适合正在做气象深度学习、遥感超分、或者对连续坐标网络建模感兴趣的人参考,不一定需要你有很深的气象背景,但最好懂一点神经网络的基本概念。
先直说这个项目要解决什么:常规的降尺度模型,无论是双线性插值、统计回归还是深度学习超分,基本都绑定在固定倍率上。比如你训练了一个输入64x64、输出256x256的模型,就只能做4倍超分;临时想输出3.7倍、9.2倍,要么重新训练,要么凑合插值。这个项目想做的,是把“尺度”从模型结构参数里解放出来,把它变成网络查询坐标的一部分。训练时你只需要在某个固定的高分辨率标签集上学习,推理时却可以任意指定目标网格密度,输出连续、平滑且物理上合理的空间场。
这套设计里,最值得展开的其实是三个关键词:任意尺度(arbitrary scale)、隐式神经网络(implicit neural networks)、混合结构(mixture)。下面我按自己的理解和工作惯例逐一拆解。
1. 项目的底层逻辑:为什么气象降尺度需要“任意尺度”
1.1 传统降尺度方案的边界在哪
我们平时拿到的大气数据,本质上是网格化的离散场,比如某个区域模式输出的是25km网格,ERA5再分析是0.25度,站点观测又是一个个零散点位。降尺度就是从粗网格往细网格“补细节”的过程。传统方案分两大流派:动力降尺度和统计降尺度。动力降尺度靠有限区域模式跑物理方程,精度高但计算成本高,跑一次几十小时很正常;统计降尺度靠历史观测建立粗网格和细网格之间的映射,速度快但非常依赖训练样本的代表性。
深度学习降尺度这几年火起来,主要是想在统计降尺度的基础上提升非线性表达能力。早期做法就是超分辨率那一套:先用卷积把粗图编码,再用转置卷积或者像素重排上采样到指定倍率。这类模型有个天然短板:上采样层是固定设计的。你用了4倍的pixelshuffle,模型就只能输出4倍分辨率。换一个倍率,网络结构就要改,训练也得重新来。业务上这个问题很现实,气象部门不可能为一个新分辨率单独训练一套模型,存储和算力都撑不住。
1.2 隐式神经网络怎么把尺度变成连续变量
隐式神经网络(INR)的思路不是从像素到像素,而是从“坐标”到“数值”。它不管输入图像多大,而是把整个场当成一个连续函数:你给它一个经纬度坐标(或者归一化后的网格坐标),它返回该位置的物理量值。这样分辨率就不再是网络结构决定的,而是由你“采样多少个坐标”决定的。想要2.5倍分辨率,就在网格上多采样一些坐标;想要7.3倍,再加密采样。网络本身完全不用动。
用INR做气象降尺度,相当于用一个MLP去拟合“大气物理场随空间坐标变化的函数”。这听起来有点疯狂,但NeRF在三维场景重建里已经证明过:纯MLP拟合坐标到物理量的映射,能够表达非常复杂的空间结构,只要给足训练数据和多尺度约束。大气温度场、气压场这类物理量天然是空间平滑的,相比图像的颜色突变,INR反而更容易收敛。
1.3 固定尺度训练,为什么不是个缺陷
这个项目标题里的“trained on fixed”一开始很容易被误解,觉得固定尺度训练等于模型只认识一个尺度,这跟任意尺度目标不是矛盾吗?其实完全不是。这里的固定尺度,指的是训练阶段的高分辨率标签相对低分辨率输入的倍率是固定的,比如都用8倍作为训练目标。模型的输入输出都发生在连续坐标空间里,推理时你可以通过调整采样密度得到任意尺度的输出。
但要注意一个实操上的坑:如果训练时只固定一种采样密度,模型容易把所有参数都过拟合到那个尺度的频谱特征上,换一个尺度推理时高频细节表现就会明显变差。我自己的做法是,标签分辨率固定为高分辨率网格,但训练时随机在低分辨率输入网格上采样坐标,等价于在“固定标签尺度”下不断变化输入网格密度,让模型学到真正的连续映射,而不是死记8倍关系。这一点在后面的实操章节我会详细展开。
2. 方案选型:为什么是“混合网络”而不是单一大MLP
2.1 单隐式网络在大气场上的瓶颈
单独一个MLP做全场映射,理论上可行,实际操作中会遇到两个问题。第一个是容量和计算量的问题:一个中等规模区域(比如几百公里范围)的0.01度网格,坐标点数超过百万级,对每一个点都要过一次MLP前向,如果网络很宽很深,推理速度会被拉到不可接受的程度。第二个是表达异质性:大气场不同区域的特征强度差异巨大,山区和海洋、晴天和锋面系统,空间平滑性完全不同。单一MLP被迫用同一组权重去拟合所有区域和所有天气状态,很容易在强梯度区域产生模糊,在均一区域产生伪振荡。
这也是我后来极力倾向混合结构的原因。混合隐式神经网络(Mixture of Implicit Neural Networks)的思路,是训练多个子网络,每个子网络负责一种空间特征或物理状态,再通过一个轻量路由网络决定每个坐标点应该更依赖哪个专家。路由的输入可以是坐标位置、局部梯度统计量、尺度因子,甚至是模式输出的辅助变量。这比单一大MLP效率高得多,也接近业务上“分区建模”的经验。
2.2 混合机制与预训练数据混合的关系
做语言模型的人可能听过regmix:data mixture as regression for language model pre-training。这个思路的核心是把“不同来源数据配比”本身当作一个回归问题来学习,而不是拍脑袋定比例。我在设计这个混合隐式网络的时候,受这个思路启发很大。降尺度任务的“数据混合”不像语言那么抽象,但性质类似:不同区域的样本、不同尺度的样本、不同背景场的样本,如果均匀混合,模型会顾此失彼;如果按区域比例抄作业,又忽略了难易样本的差异。
所以我在路由网络里不只是让专家做空间分区,还会把“区域训练样本的统计分布”作为路由输入的一部分。这就把数据混合的问题变成了一个可微的路由权重回归:哪个专家在哪个区域、哪个梯度强度下更可靠,由训练过程自己摸索。相当于每一个输出坐标点都有一个动态的专家分配,而不是简单地把区域切成几块。
2.3 整体架构与链路设计
我最终采用的方案分三块:坐标编码器、若干专家网络、路由网络。先说坐标编码器,输入的是归一化坐标,通常会先用傅里叶特征映射到高频空间,让MLP能表达细节丰富的温度场结构。专家网络是几组结构相同但参数独立的MLP,每个专家内部有自己的频段偏好,之所以能做到这一点,是靠初始化和训练数据分配的自然分化。路由网络则比较轻,输入坐标和对应的低分辨率上下文特征,输出每个专家的权重。最后加权求和得到该坐标点的预测值。
训练时,低分辨率输入场会先通过一个小型编码器得到整体上下文特征,这个特征会和坐标一起送到路由和专家网络。这样专家就不只是看坐标,还能感知当前场的局部状态,比如知道当前正处于强锋区还是均一洋面,从而调整选择策略。推理时,外部只需提供目标网格的坐标序列和低分辨率输入场,就能批量产出任意尺度的预测网格。
3. 实操过程与核心环节实现
3.1 数据准备:从粗网格到连续映射的三级处理
数据是一切的前提。气象降尺度项目里,我通常用ERA5再分析做标签,因为它的空间覆盖和物理一致性都比较好。这里需要加载的变量包括近地面温度、湿度、风场、海平面气压等。不要直接拿原始网格开训,要先做三级处理。第一级是裁patch,把研究区域切成若干重合的子块,既能减少存储压力,也能让每个专家更容易学局部特征。第二级是去趋势,把气候态均值先减掉,让网络集中学“异常场”的降尺度规律,否则大部分网络容量都用在背气候态上了。第三级是坐标归一化,把所有坐标映射到[-1,1]区间,这对MLP的收敛非常关键。
要特别注意一个细节:坐标和变量值的配对关系必须严格保留。你在空间域裁剪patch,变量场对应裁下来,坐标网格也要同步裁。我早期犯过的错误是把坐标统一归一化成整块大网格的标准坐标,导致裁出来的每个patch在模型看来都挤在同一个坐标区间,空间信息严重失真。正确的做法是每个patch用自己实际的范围做归一化。
3.2 关键模块:坐标编码、专家与路由的PyTorch实现
光说架构太悬浮,这里贴一段我简化后的核心代码结构,可以当模板直接用。
import torch import torch.nn as nn import math class FourierFeatureEncoder(nn.Module): def __init__(self, in_dim=2, num_bands=64, scale=10.0): super().__init__() self.num_bands = num_bands self.scale = scale def forward(self, coords): # coords: [B, N, 2] normalized to [-1, 1] freq = 2 ** torch.linspace(0, self.num_bands - 1, self.num_bands) freq = freq * math.pi * self.scale / (2 ** (self.num_bands - 1)) freq = freq.to(coords.device).view(1, 1, -1) proj = coords.unsqueeze(-1) * freq # [B, N, 2, B] feat = torch.cat( [torch.sin(proj), torch.cos(proj)], dim=-1 ).flatten(-2) # [B, N, 2*2*B] return feat class ExpertMLP(nn.Module): def __init__(self, in_dim, hidden_dim=256, out_dim=1, num_layers=4): super().__init__() layers = [nn.Linear(in_dim, hidden_dim), nn.ReLU()] for _ in range(num_layers - 1): layers += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU()] layers += [nn.Linear(hidden_dim, out_dim)] self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x) class MixtureImplicitNetwork(nn.Module): def __init__(self, num_experts=4, coord_dim=256, ctx_dim=16): super().__init__() self.encoder = FourierFeatureEncoder() self.experts = nn.ModuleList( [ExpertMLP(coord_dim + ctx_dim) for _ in range(num_experts)] ) self.router = nn.Sequential( nn.Linear(coord_dim + ctx_dim, 64), nn.ReLU(), nn.Linear(64, num_experts) ) def forward(self, coords, context): feat = self.encoder(coords) # [B, N, coord_dim] ctx = context.unsqueeze(1).expand(-1, coords.size(1), -1) x = torch.cat([feat, ctx], dim=-1) logits = self.router(x) weights = torch.softmax(logits, dim=-1) # [B, N, E] pred = 0 for i, expert in enumerate(self.experts): pred = pred + weights[..., i:i+1] * expert(x) return pred这个结构里最值得琢磨的是context变量的设计。我这里的上下文特征不是全局平均那种粗糙的东西,而是低分辨率输入场经过一个小卷积编码器输出的局部语义向量,代表这个坐标点附近的大气状态。你可以把它理解成给路由和专家递了一张“小抄”,告诉它们当前气象条件更接近锋面、暴雨还是晴空。有了上下文,混合机制才不会变成纯粹的空间分区。
3.3 训练设置与损失函数的取舍
训练损失我采用L1为主、SSIM为辅的组合。L1对异常值的惩罚比L2温和很多,气象场尤其是降水这类重尾分布变量,用L2容易把预测拉向均值,靠近零值区域会糊成一片。SSIM可以帮助模型在结构上保持空间相关性,代价是训练更慢一点。如果业务上特别看重极值区域,还可以额外加一个分位数损失项,但我个人经验是先把L1和SSIM调稳,再考虑极值约束。
网络输入的低分辨率场是通过对高分辨率标签做固定倍率下采样得到的。我通常设4倍或8倍。最关键的trick是训练时不要只按固定网格采样坐标。每次迭代随机在标签网格上选取一部分坐标点,而不是把所有像素都用上,这样一方面节省显存,另一方面让模型在“部分观测”下依然能够重建完整连续场,隐式网络的泛化性会好很多。
学习率方面,我用AdamW,初始1e-4,配合cosine退火。专家网络的参数更新会比路由略快,如果发现路由权重长期不分化,就把路由的学习率降低一个数量级。混合网络的两个部分需要平衡,不然后面会碰到路由坍缩,具体表现和解决方法我放在下一章。
3.4 任意尺度推理:一次前向输出任意分辨率网格
推理的时候,过程非常简单。我先确定目标网格的行列数,比如你想从0.25度降尺度到0.05度,那就按0.05度间隔生成经纬度坐标,归一化之后一次性丢给模型。模型输出的是每个坐标对应的物理量值,reshape成目标网格,就是你要的任意尺度结果。面密度完全由你提供的坐标数量决定,跟模型结构毫无关系。
这里有个性能优化点:百万级坐标直接过一个宽MLP,显存会瞬间爆炸。我的做法是分块推理,每次处理4096或者8192个坐标,全部输出后拼回完整网格。实测下来,和批量反正是一样的效果,但显存占用能压到可接受范围。另外一个优化是,如果某些专家权重在所有坐标上都很低,可以在推理阶段跳过它们,这需要提前统计路由权重,属于工程上的锦上添花。
4. 走了不知道多少弯路:常见问题与排查实录
4.1 固定尺度过拟合与尺度泄漏
固定尺度训练最典型的坑,就是模型会对“固定倍率”这个隐含条件产生依赖。直观感受是:在训练那个倍率上测试效果很好,SSIM高得离谱;换一个推理倍率后,硬件细节突然变得模糊,甚至出现条状伪影。起初我怀疑是坐标编码的问题,后来逐步定位到是模型把采样密度当成了隐身特征。解决方法是scope augmentation:训练时让低分辨率输入和高分辨率标签的对应关系在一个区间内随机浮动。比如核心倍率是8倍,我可以按0.5到1.5的随机系数扰动实际采样密度,让模型无法偷懒固定映射。
另一个容易忽略的点是坐标精度。如果你在GPU上用float16存储归一化坐标,高频傅里叶特征的高频分量很容易被截断,导致输出出现“网格状”纹路。我自己排查了一整天才发现是自动混合精度挖的坑。建议关键坐标计算保持float32,或者把傅里叶特征的计算提到混合精度开关之前。
4.2 混合网络的路由坍缩
混合结构里最经典的问题,莫过于路由权重退化成one-hot:不管输入的坐标和上下文是什么,模型永远只选某个专家,其他专家形同虚设。出现这个情况时,模型效果甚至可能比单个MLP还差,因为参数浪费了。根因通常是负载不均衡的样本分布,某类天气状态的样本量太大,路由发现全压给一个专家就能把loss压下来。
我的处理手段有三种:第一,给路由加负载均衡正则项,约束每个专家平均接收的样本数不要偏离太远;第二,给路由logits加噪声,训练初期让专家分配更多样化,随后噪声逐渐退火;第三,对样本按天气类型或区域做重采样,从数据层面缓解不平衡。这三种方法同时用效果最好,单独用任何一种都容易复发。
4.3 边界不连续和伪影
隐式网络做任意尺度推理时,如果目标区域超出了训练patch的坐标范围,边界处会出现明显的不连续。这本质上是外推问题:MLP只能表达训练分布内的映射,你让它预测训练区域外十公里的坐标,它找不到参考样本,就会输出一个无意义的常数。项目里我的解决方式是把训练patch设计成有重叠区域,并且在推理时只取patch中心部分的预测结果,边缘交给相邻patch的模型输出做加权融合。
气象场上还有一个特殊的边界问题:区域模式边界的物理通量是不守恒的。如果只随机裁patch,模型学不到大尺度背景环境的约束,降尺度结果可能局部形态好看,但区域平均能量漂移。一个补救措施是在输入上下文里加一个大尺度环境场特征,让模型知道当前patch在整个大区域的相位位置。这么做之后,边界伪影明显减少。
4.4 评估指标的误导性
只看PSNR和SSIM评估降尺度质量,容易得出乐观结论。特别是SSIM,它对整体亮度偏移特别不敏感。气象业务真正关心的是:极值中心的位置准不准、区域总量守恒不守恒、空间频谱是否合理。所以我的评估方案是多指标的:PSNR和SSIM保留供和同行对比;新增RMSE分尺度统计,把0.05度、0.1度、0.25度分开算;对变量场做功率谱分析,看高频能量有没有异常堆积;如果是降水,还会算CRPS来评估概率分布的匹配程度。
实测发现,论文里那些PSNR涨了0.3dB的模型,在CRPS上可能反而变差。这说明单纯向着像素指标优化,会牺牲分布合理性。做气象深度学习,不能只看图像质量指标,这事我强调再多也不嫌多。
| 问题现象 | 可能原因 | 优先级排查方向 |
|---|---|---|
| 固定倍率效果好,其他倍率崩 | 固定尺度过拟合,坐标采样密度被当成隐含变量 | 增加scale augmentation并检查坐标精度 |
| 路由长期选同一专家 | 样本分布不平衡或负载正则缺失 | 加路由噪声和负载均衡正则 |
| 推理出现网格状纹路 | float16坐标截断或傅里叶基频设置过于极端 | 坐标特征用float32计算,降低band数量 |
| 区域边界出现跳变 | 训练patch坐标范围外推,上下文背景丢失 | 用重叠patch预测并在边缘做加权融合 |
| 降尺度场整体偏平滑 | 高频损失权重不足,或者损失函数L2占比过重 | 调整L1与频域损失、合理使用感知损失 |
4.5 显存优化:让百万级坐标的推理落地
气象网格动辄上千万点,直接全体坐标前向,再好的GPU也会被显存吃穿。我上面提到分块推理,这是最实用的方案。具体的块大小可以依照激活函数所占显存来估算,我一般从8192开始试,如果显卡负载不错就逐步加大,直到刚好不爆显存为止。还有一种思路是把专家网络和路由网络拆开,先批量算路由权重,再按专家分组聚合样本,减少重复计算。
模型推理速度在业务上也很关键。如果再分析数据更新一轮只需要六小时,但降尺度模型要跑十小时,那就没有实用价值。我的实测配置下,单个专家MLP四层256维,8个专家,分块推理在A100上处理一个中等区域的0.01度网格大概在几十秒量级,完全满足业务时效。如果换更小的模型或者减少专家数量,还能进一步压缩,只是精度会有轻微下降。
5. 这些经验还能往哪些方向延伸
5.1 引入静态协变量,让专家分工更明确
大气降尺度不只有网格坐标这一种输入,地形高度、土地利用类型、距离海岸线距离等静态变量,对局地气候有决定性影响。把这些作为辅助协变量拼接到坐标特征里,能有效降低路由网络的负担,因为它不需要仅靠坐标去隐式推断地形位置。我在温度场的实验里加入地形高度之后,山区站点的RMSE下降了将近两成。这个改动很便宜,就是把协变量归一化后拼进编码器输入,强烈建议后续直接纳入。
5.2 多变量一致性与物理约束
目前项目集中在单一变量场,比如只做温度或者只做降水。业务上最怕的是温度降尺度得很好,但湿度场和风场跟它对不上,导致后续物理量诊断出错。下一步可以考虑多变量联合降尺度,让每个专家输出多个变量。代价是专家网络的输出头变宽,训练难度上升。另一个方向是加入简单的物理约束损失,比如温度随高度的递减率、水汽连续性之类的弱约束,让模型输出不违反基本大气规律,这在样本外的稳定性上会有帮助。
5.3 数据混合的思路还可以继续深挖
regmix那个思路给我的启发不止停留在路由设计上。它本质上是说“数据的构成方式应该被模型自动学习,而不是被人类预设”。在混合隐式网络里,数据构成对应的是训练patch采样策略。实际上每个patch的天气状态、地形复杂度、目标难度完全不一样,采样时如果对难度大的区域过采样,模型在小概率环境下的表现会提升。但手写难度度量又很主观,一个可行的做法是,用当前模型的预测误差作为采样权重,每次训练迭代动态调整各patch被抽到的概率。
这个做法实现起来很简单,只需要维护一个误差指数加权列表,但效果非常明显。我在第二次迭代实验时用了这个策略,极值区域的预测精度提升很快,而且没有明显伤害均一区域的性能。调参上要注意平滑系数,误差更新太快会导致训练后期采样分布抖动过大。
最后说说我的整体感受。任意尺度降尺度这个方向,视野上比传统固定倍率超分大得多,一次训练、任意分辨率输出,这种灵活性对业务系统来说价值极高。混合隐式神经网络让我可以用一群相对小的网络去分工覆盖大气场里的不同物理行为,比堆一个大网络更有解释性,训练可控性也更好。固定尺度训练看起来像一个约束,但配合坐标连续采样之后,反而成为模型稳定的一个锚点。做这类项目,最大的心得就是别把尺度当作离散类别去学,要让模型在连续坐标里理解场结构;也别指望单个万能网络吃遍所有区域,混合和路由是一个工程上更现实的选择。后面如果条件允许,我很想把多变量和物理约束这两个方向继续推进,让这套方法真正落到气象业务里,而不只是停留在实验指标上。