简介傅里叶神经算子FNO是解决参数化偏微分方程的一种新型深度学习方法这份开源压缩包正是其完整实现适用于物理模拟、流体力学与工程计算等领域。包内共含有二十九个文件以二十一个Python源码为主另附六个m文件、一个Markdown说明文档和一个许可证文件压缩后大小仅约六十KB结构紧凑但复现路径清晰。源码覆盖数据预处理、快速傅里叶变换、频域卷积层、逆傅里叶变换及损失函数等全流程并有一维、二维、三维和时间相关模型变体还自带Darcy、Navier-Stokes、Burgers等经典偏微分方程的数据生成样例便于系统理解模型机制、部署技巧和迁移应用。目前已有四百一十四人学习该资源适合具备一定深度学习基础的研究者作为入门手写实现与二次开发的起点。1. 为什么 PDE 代理模型绕不开 FNO一个压缩包背后的技术含量做流体仿真或者计算电磁的同学大概都有这种体验传统求解器算一次达西流或者二维涡度方程网格细一点就是分钟级甚至小时级而参数扫描动辄几百组工况根本算不动。fourier_neural_operator这个压缩包解压之后就是一套基于傅里叶神经算子FNO的 PDE 代理模型实现。它的核心思路是把偏微分方程的求解过程学成一个函数空间到函数空间的映射训练完成后单次推理在毫秒级完成比传统迭代求解快两到三个数量级。这个压缩包能帮你解决的问题很直接不用动求解器内核把数值模拟的输入输出当监督数据训练一个傅里叶层网络替换掉耗时的迭代求解过程。适合手里有仿真数据、想给求解过程提速但不想从零手写神经算子的工程师和研究生。2. 先看懂 FNO 在做什么傅里叶层与算子学习的核心原理2.1 神经算子与传统神经网络的本质区别从「拟合函数」到「学习映射」传统 CNN 或 MLP 解决的是函数拟合问题输入是一个固定维度的向量或者网格输出是另一个固定维度的向量或网格。例如图像分类输入 224×224×3输出 1000 类概率输入输出的形状在训练前就锁死了。这种模式有一个隐含限制——模型学到的变换与离散化的分辨率绑定换一套网格或者加密采样模型往往要重训。FNO 走的是另一条路。它学习的不是函数而是算子。算子在泛函分析的语境下是「函数的函数」输入一个函数例如 PDE 的系数场 a(x)输出另一个函数例如速度场 u(x)。因为输入输出都是函数理论上它们在任意分辨率下都有定义模型训练好之后可以对没见过的新网格做推理这就是 FNO 常被提起的「网格无关性」mesh-invariance。用公式来表达更直观。设输入函数为 a输出函数为 uFNO 要逼近的是映射 GG(a) u。实际计算时输入被离散成网格上的采样值但网络结构本身不依赖网格尺寸。每一层的作用是更新函数表示v_{t1}(x) σ(W v_t(x) (K v_t)(x))其中 W 是逐点线性变换K 是一个全局积分算子也就是下一节要展开的傅里叶层。这个形式在数学上对应积分方程迭代求解的结构所以 FNO 对 PDE 问题尤其贴合。2.2 Fourier Layer 的计算流程FFT、谱域线性变换与参数共享Fourier Layer 是 FNO 的核心组件名字听着玄学拆开看其实只有四步升维、FFT、谱域线性变换、IFFT。我用一段伪代码风格的结构来说明实际 PyTorch 实现也基本是这个骨架。# 以二维 FNO 为例输入张量形状为 (batch, in_channels, H, W) def fourier_layer(v, modes1, modes2, weight_r): # 1. 残差分支逐点 1x1 卷积保持空间形状不变 v_res pointwise_conv(v) # (batch, hidden, H, W) # 2. 对空间维度做 FFT得到频域表示 v_hat torch.fft.rfft2(v) # (batch, hidden, H, W//21) # 3. 截断高频只保留前 modes1, modes2 个低频模态 v_hat_trunc v_hat[:, :, :modes1, :modes2] # 4. 谱域线性变换复权重矩阵作用于通道维 v_hat_out torch.einsum(bchw,cod-bohw, v_hat_trunc, weight_r) # 5. 补零回原尺寸IFFT 回到物理空间 v_out torch.fft.irfft2(pad_zero(v_hat_out), s(H, W)) # 6. 加上残差分支过激活函数 return torch.relu(v_res v_out)代码里的关键点有两个。第一是modes1和modes2这是谱截断的模态数量通常取 8 到 24 之间。它的物理含义是只保留频域中最主要的低频分量把高频部分直接扔掉再通过可学习的复数矩阵weight_r对保留的模态做线性组合。这个操作对 PDE 解的光滑场特别合理因为大多数物理场的能量集中在低频段。第二是torch.einsum里的weight_r张量形状是(in_channels, out_channels, modes1, modes2)它本质上定义了一个频域上的卷积核参数数量远小于物理空间的全局卷积核这是 FNO 高效的原因。与普通卷积不同傅里叶层的卷积核在频域是全局的。物理空间的一个小卷积核只能看到局部邻域而频域的一个模态对应整个空间域的全局模式所以一层傅里叶层就能捕获全域依赖不需要堆几十层来扩大感受野。2.3 为什么选傅里叶基而非卷积全局感受野与网格无关性既然要做全局建模理论上用 Transformer 的自注意力也能做到但 FNO 选择了傅里叶基原因是计算复杂度。标准自注意力的复杂度是 O(N²)N 是网格点数二维 128×128 网格的自注意力直接爆显存。而 FFT 的复杂度是 O(N log N)加上模态截断实际算得很快。傅里叶基还有两个附加优势。一是天然的抗噪正则化截断高频等价于对解做了一次平滑滤波训练时不容易被数据里的数值噪声带偏。二是参数效率一个modes1 × modes2的频域权重矩阵共享了整个物理空间的所有位置这跟 CNN 的权重共享思想一致但覆盖范围是全局的所以同等参数量下表达能力更强。不过傅里叶基也有边界。如果解本身包含强间断比如激波问题低频截断会造成 Gibbs 振荡解附近会出现明显的波纹。这类问题要么把模态数调大、在损失里加高频惩罚要么改用带小波基或自适应基的变体。选型的时候先确认你的解场是光滑的还是带间断的再做决定。3. 把 fourier_neural_operator 跑起来环境搭建与最小训练示例3.1 解压与目录结构拿到压缩包后先看什么解压这个压缩包之后先别急着跑训练花两分钟把目录结构过一遍。常见的组织方式大致是data/、models/、train/、utilities/几个目录外加一个 README不同版本会有差异但核心模块基本都在。unzip fourier_neural_operator-master.zip cd fourier_neural_operator-master ls -la # 典型结构 # README.md # data/ # 数据生成与加载脚本 # models/ # FNO 模型定义fno.py / fno_2d.py 等 # train/ # 训练入口脚本 # utilities/ # 归一化、损失函数、可视化工具先看 README确认这个版本支持的 PyTorch 版本和 Python 版本再进models/看一下模型目录的forward函数确认输入张量的维度约定。我见过不少同学跳过这步直接跑训练结果维度对不上报错回头排查半天才发现是数据加载器的输出排布和模型预期不一致。这个包的数据加载器一般会把输入输出拼成一个字典返回字段名可能是x和y也可能是coeff和sol训练脚本里会有对应代码先读通一个 batch 的数据流再动手。3.2 环境依赖与安装一版不会翻车的组合FNO 的依赖不复杂核心是 PyTorch、NumPy、SciPy数据后处理可能用到 h5py。最稳妥的方式是新建虚拟环境不要直接往 base 环境里装。# Python 3.8~3.10 均可PyTorch 2.x 或 1.13 都行 python -m venv fno_env source fno_env/bin/activate pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy scipy h5py matplotlib tqdm装 PyTorch 的时候注意 CUDA 版本要和本机驱动匹配nvidia-smi看一下驱动支持的 CUDA 版本再选对应的--index-url。如果没有 GPUCPU 也能跑通小规模数据只是训练速度会慢一个数量级建议先用 64×64 的小网格验证流程再上大网格。装完之后从models/里 import 模型做个 2 秒的 smoke test确认 FFT 相关算子能正常执行。3.3 最小训练命令把默认配置跑出第一个 loss 曲线环境装好之后先用默认配置把训练流程跑通。大部分版本的训练入口都是一个train_XXX.py脚本接收命令行参数控制主要超参数。# 以二维达西流为例跑 20 个 epoch 验证流程 python train/train_darcy.py \ --batch_size 16 \ --epochs 20 \ --learning_rate 0.001 \ --modes 12 \ --width 32 \ --data_path ./data/darcy \ --log_dir ./logs/darcy_test训练开始后观察 loss 曲线的下降趋势。正常情况前 5 个 epoch loss 会下降一到两个数量级之后进入缓慢下降。如果 loss 完全不动先检查归一化统计量是不是计算错了如果 loss 下降但验证误差不降大概率是过拟合需要加数据量或调小width。# 让训练输出更可控固定随机种子保证可复现 python train/train_darcy.py --seed 42 --amp --log_dir ./logs/darcy_test_v2--amp开启混合精度训练傅里叶层的 FFT 在 float16 下能显著省显存但注意torch.fft在 float16 下某些 CUDA 设备上会有精度损失如果训练不稳定就关掉。3.4 训练脚本核心参数epochs、batch_size、learning_rate 怎么改训练脚本里最常见的几个参数改之前先想清楚它们之间的关系。batch_size影响的是梯度的平滑程度和显存占用FNO 的 FFT 计算在 batch 维是并行展开的batch 从 16 增到 32显存增长接近线性如果显存不够优先减 batch 而不是减网格分辨率。learning_rate用 0.001 起步基本不会翻车但训练后期最好加上 StepLR 或者 CosineAnnealingLRFNO 的损失曲面在高频分量上比较敏感固定学习率容易在后期震荡。# 典型训练循环核心逻辑来自常见实现关键代码已标注 optimizer torch.optim.Adam(model.parameters(), lrargs.learning_rate, weight_decay1e-4) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size100, gamma0.5) for epoch in range(args.epochs): model.train() train_loss 0.0 for batch in train_loader: # 数据形状: (batch, 1, H, W)归一化已在数据加载器完成 x, y batch[x].to(device), batch[y].to(device) optimizer.zero_grad() pred model(x) loss rel_l2_loss(pred, y) # 相对 L2 损失见第 6 章 loss.backward() optimizer.step() train_loss loss.item() * x.size(0) scheduler.step() print(fEpoch {epoch}: train_loss {train_loss / len(train_loader.dataset):.6f})weight_decay1e-4是经验值用于抑制频域权重过大导致的过拟合。StepLR的step_size100和gamma0.5意味着每 100 个 epoch 学习率减半这个配置是按默认 500 epoch 设计的如果你的epochs只有 50可以把step_size改到 20。训练循环里那个rel_l2_loss它计算的是预测场相对真实场的 L2 范数误差。注意这个最小示例跑通的意义在于验证环境和数据链路不要在默认参数上花太多时间刷精度。默认参数往往针对特定数据集调过换数据后需要重新调。4. 数据准备与模型配置从达西流到 Navier-Stokes 的实操参数4.1 数据维度约定Batch、Channel、Spatial 的排布顺序FNO 模型对输入张量的排布顺序是敏感的。PyTorch 默认的排布是(batch, channel, H, W)FNO 的二维模型也遵循这个约定但一维 FNO 稍有不同输入是(batch, channel, L)其中 L 是空间采样点数。很多人在做时序 PDE 问题时把时间维塞进了 channel 或 batch导致模型效果出现问题。# 数据加载器中推荐的张量构造方式 # 输入场 a: 网格上的系数场形状 (N, H, W) # 输出场 u: 对应的解场形状 (N, H, W) import numpy as np import torch a np.load(coeff_field.npy) # (N, H, W) u np.load(sol_field.npy) # (N, H, W) # 在 channel 维插入一个维度变成 (N, 1, H, W) a_tensor torch.from_numpy(a).float().unsqueeze(1) u_tensor torch.from_numpy(u).float().unsqueeze(1) # 如果输入有多个物理量比如 NS 方程的初始速度场和压力场拼在 channel 维 # a_multi torch.cat([u0_tensor, v0_tensor], dim1) # (N, 2, H, W)数据准备的核心是把你的物理场打包成「输入场 输出场」的监督对。输入场是 PDE 的系数或者初始条件输出场是求解器算出来的最终解。如果原始数据是结构化网格上的文本或 VTK用 SciPy 的griddata插值到规则网格再做归一化。4.2 自制数据的输入输出配对监督标签怎么构造FNO 的训练数据说起来简单——一对一对的输入输出场——但实际操作中配对的细节决定模型能不能学出来。第一是空间对齐输入场的网格和输出场的网格必须在同一套坐标系下如果是内嵌边界或者不规则域先把边界外的值 mask 掉不然模型会去学边界外的无效信息。第二是归一化输入和输出分别做标准化统计量只从训练集计算验证和测试集沿用训练集的均值和方差这一步很多人做反了导致验证误差虚高。# 归一化统计量只从训练集计算保存下来供推理时复用 from sklearn.preprocessing import StandardScaler scaler_x StandardScaler() scaler_y StandardScaler() # 训练集形状展平成 (N, H*W) 再 fit避免逐通道统计 N, C, H, W a_tensor_train.shape scaler_x.fit(a_tensor_train.reshape(N, C*H*W).numpy()) scaler_y.fit(u_tensor_train.reshape(N, C*H*W).numpy()) def normalize(x, scaler): # 恢复成 (N, C, H, W) return torch.from_numpy(scaler.transform(x.reshape(N, -1))).reshape(N, C, H, W)我一般会提醒PDE 代理模型的输入场往往有多个物理量比如二维 NS 方程需要初始时刻的 u、v 速度场外加压力场三个通道的数值量纲不同绝对数值差异可能差好几个数量级。这时候逐通道做标准化比全局标准化效果好否则量纲大的通道会主导损失函数。4.3 modes 和 width 的设置逻辑谱截断是分辨率还是特征宽度modes是 FNO 的特有超参数刚接触的人容易把它误当成卷积核大小——其实它的含义完全不同。modes控制的是 FFT 之后保留多少个低频模态数。以二维情况为例modes112, modes212的意思是频域只保留前 12 行的低频系数其余全部丢弃。这个参数决定了模型能捕获的最低频率分辨率设置太小解的高频细节学不出来结果会偏平滑设置太大除了增加计算量还可能把数据里的数值噪声当成信号学进去。width则是通道数对应标准 Transformer 里的 hidden dimension 或 CNN 里的 filter 数量。width控制每个空间位置的特征表达能力。经验上modes和width需要一起调width增大后频域权重矩阵的通道维变宽同样数量的模态能组合出更复杂的频域变换所以width大时modes可以适当减小反过来modes增大后每个模态的通道变换也需要更多参数来利用这些信息width不变的话提升有限。调参时我习惯按「先固定一个、调另一个」的方式做网格搜索。第一轮固定width32把modes从 4 到 20 扫一遍画误差曲线找拐点第二轮固定最优modes把width从 16 到 128 扫一遍。这个顺序能避免两个参数同时变动时分不清是谁的贡献。4.4 不同 PDE 场景的推荐初始配置一页表格不同 PDE 问题的数据特性和难度差别很大下面这组初始配置是我多次实验后觉得起步比较稳的参数组合可以作为调参起点而不是最终答案。问题类型数据规模modeswidthbatch_sizelearning_rate备注一维 Burger 方程1000 条轨迹1632320.001时间维用 FNO 的 Lifting 层处理达西流二维1000 组系数-解场12×1232160.001光滑解场容易收敛Navier-Stokes 二维1000 组初始-终态16×166480.0005通道多显存占用高batch 要小带间断的守恒律500 组激波解24×246480.0003模态要放大否则 Gibbs 振荡严重NS 问题显存吃紧的原因不只是width64带来的参数量还包括 Fourier Layer 中 FFT 的中间变量每个样本在频域的 shape 是(64, 64, H, W/21)的复数张量128×128 网格下这个中间变量就有几十 MB。如果显存不够优先把batch_size降到 4其次是降网格分辨率到 96×96 而不是降width因为降width对精度的影响比降分辨率更直接。注意这个表只适合 64×64 到 128×128 的网格。网格到了 256×256FFT 的计算量和显存占用会显著上升通常需要改成分布式训练或者用 FNO 的分块版本。5. FNO 的 5 个高频踩坑与排查思路5.1 训练 loss 下降缓慢甚至不降先查 normalization 和 weight_init现象loss 在 1.0 附近徘徊几十个 epoch 都不动或者下降速度极慢像是有个看不见的力量在拖着模型。原因最常见的是输入和输出场的数值范围没有归一化。如果输出场的量级是 1e4相对 L2 损失的初始值就会非常大梯度计算时被大数值主导Adam 的动量项被带偏导致收敛极慢。另一个原因是模型最后一层的权重初始化太大输出的初始预测值飞出合理范围。解决检查数据加载器里的标准化逻辑确认scaler_y的均值方差是合理的再把模型最后一层的权重手动缩小。通常把最后一层线性层的bias初始化为训练集输出的均值效果立竿见影。# 常见实现中最后一层线性层的初始化修正 for layer in model.output_layer.children(): if isinstance(layer, torch.nn.Linear): # 用训练集的输出均值初始化 bias加速收敛 layer.bias.data.fill_(y_mean.item()) # 权重用小随机数初始化避免初始预测过大 layer.weight.data.mul_(0.01)5.2 推理误差突然变大输入分辨率变了但 modes 没跟着调现象训练时验证误差 1% 左右换到高分辨率输入做推理时误差飙升到 20% 以上。原因FNO 的「网格无关性」是有条件的。如果输入空间分辨率从 64×64 变成 128×128但modes还是 12那么频域被截断的物理频率范围就变了——同样的模态编号在更高分辨率下对应更高的物理频率模型训练时没见过这些频率的输入模式。另一个常见问题是推理时输入数值范围没变但归一化统计量来自训练集如果测试数据分布有偏移误差也会变大。解决推理时保持输入分辨率接近训练数据的分辨率范围如果应用场景必须换分辨率重新用新分辨率生成少量数据做 fine-tune把modes按分辨率倍数相应调大。我用过的一个经验值分辨率翻一倍modes大约增加 1.5 倍可以保持频域覆盖范围近似。5.3 显存爆掉width 增大后 FFT 的显存放大效应现象训练能跑把width从 32 调到 64 之后直接 OOM而且不是渐进地变慢是突然爆掉。原因Fourier Layer 的 FFT 中间表示在频域是复数张量实部虚部各占一份显存实际占用是物理空间张量的 2 倍再乘上width的通道数。width翻倍FFT 中间变量的显存占用约等于翻 4 倍通道维乘以频域复数两倍远比想象中夸张。解决先用torch.cuda.max_memory_allocated()跑一个 batch 测实际峰值显存再决定降什么。优先降batch_size其次降网格分辨率最后才降width。另外开启torch.backends.cudnn.benchmark True能小幅降低 FFT 的显存碎片。# 训练前显存自检打印一个 batch 的峰值显存占用 model model.to(device) x torch.randn(args.batch_size, 1, 128, 128).to(device) _ model(x) print(fPeak mem: {torch.cuda.max_memory_allocated() / 1e6:.1f} MB)5.4 结果出现明显高频噪声Padding 与 Gibbs 现象现象预测的解场整体形状是对的但边界附近或者间断附近出现密集的波纹振荡肉眼可见的不光滑。原因两个来源。一是 FFT 的周期性假设torch.fft.rfft2默认把输入当作周期信号处理如果真实物理场的边界不满足周期性傅里叶层在边界处会强行拟合出周期性跳跃造成伪振荡。二是模态截断导致的 Gibbs 现象对带间断的场截断高频后间断点附近必然出现波纹。解决训练前把数据在空间维做 Padding——我倾向于训练时用更大的网格比如原始 64×64 外边 pad 成 80×80推理时裁回 64×64。这能明显缓解边界伪振荡。对 Gibbs 现象则要从数据或损失入手在损失函数里加一项空间梯度惩罚项强制预测场的梯度在间断点之外保持平滑。# 带梯度惩罚的损失平滑预测场的非间断区域 def pde_loss(pred, target, lambda_grad0.1): rel_l2 torch.norm(pred - target) / torch.norm(target) # 计算预测场的空间梯度惩罚过大的梯度 dx torch.diff(pred, dim-1) dy torch.diff(pred, dim-2) grad_penalty torch.mean(dx**2 dy**2) return rel_l2 lambda_grad * grad_penalty5.5 复现不了论文数字数据版本、归一化统计量和训练轮次对齐现象照着默认配置训练验证误差就是比 README 里写的数字高一大截怀疑自己代码写错了。原因大概率不是代码问题而是数据版本或训练流程的细节差异。很多公开数据的生成脚本包含随机种子、网格分辨率、求解器容差等参数数据生成时这些参数不同数据本身的难度就不同。另外归一化统计量如果用了全数据集而不是只用训练集验证集的信息泄漏会虚高测试精度——模型在「看见过」验证集分布的情况下训练你复现出的误差自然比真实的低。解决复现时注意三点——确认数据生成脚本的随机种子是否一致归一化统计量严格从训练集计算训练轮次跑满默认值通常是 500 epoch不要用 50 epoch 的结果去对比 500 epoch 的宣称精度。我从这个坑里学到的教训是先把 README 里的复现步骤逐条打钩再开始改任何东西否则后面排查都不知道往哪个方向找。6. 让 FNO 真正能用误差诊断与超参数调优技巧6.1 一组误差诊断工具相对 L2 误差与逐点误差热图训练完之后别只看一个测试集平均误差就结束。我做误差诊断的习惯是四件套相对 L2 误差、逐点绝对误差热图、预测与真实场的剖面线对比、频谱差异图。前两件必做后两件视情况。# 推理与误差诊断脚本 model.eval() with torch.no_grad(): pred model(x_test.to(device)).cpu() # 相对 L2 误差整体精度的核心指标 rel_l2 torch.norm(pred - y_test) / torch.norm(y_test) # 逐点绝对误差找出误差集中的空间位置 abs_err (pred - y_test).abs().squeeze().numpy() # 可视化绘制误差热图确认是否有空间分布偏向 import matplotlib.pyplot as plt plt.imshow(abs_err, cmaphot), plt.colorbar() plt.title(fAbsolute error field, rel_L2{rel_l2:.4f}) plt.savefig(error_heatmap.png, dpi150)误差热图能告诉我模型在哪类区域学得不好。常见模式有两种误差集中在边界或角落说明是周期性假设的边界伪影误差集中在场的高曲率区域说明modes不够高频细节学丢了。针对前者加 padding针对后者调大modes。6.2 超参数调优顺序先 modes 后 width 再训练策略超参数调优不要同时动多个旋钮我习惯按固定顺序来先扫modes固定width再以最优modes扫width最后调训练策略学习率、轮次、归一化方式。为什么这个顺序因为modes直接决定模型能看到的频率范围是 FNO 的「天花板」参数它设低了再怎么调width都是白费width决定模型利用这些模态的表达能力训练策略最后调是因为模型容量对了之后训练相关的调参才有意义。我还有一个经验是做完一个实验后把配置和误差记下来否则两三天之后再跑就忘了当初哪组参数最好。我习惯用一个简单的表格记录实验序号、modes、width、epochs、rel_l2、备注时间长了这就是最有价值的调参手册。FNO 的训练在数据量足够时其实相对稳定真正拉开差距的往往不是某个魔法参数而是对问题结构的理解——你的解场是光滑的还是比较尖锐的边界条件是不是周期的训练数据覆盖了多大的参数空间。想清楚这几个问题比盲调参数值钱得多希望帮到你。本文还有配套的精品资源点击获取
