傅里叶神经算子FNO实践指南：用深度学习替代传统PDE求解器

简介：傅里叶神经算子（FNO）的完整实现源码包，专为研究物理信息机器学习、探索偏微分方程与深度学习结合方法的开发者设计，也可作为工程仿真团队的扩展基准。FNO借助傅里叶变换捕捉空间频率特征，通过频域卷积与非线性激活直接学习PDE解的算子，可高效预测不同参数条件下的解，广泛应用于流体力学、电磁学、热传导等物理问题。压缩包共二十九个文件，主要包含二十一个代码脚本、六个数学计算脚本、一个说明文档及一份许可，其中核心脚本负责模型定义与训练，辅助脚本用于数值分析与可视化，整体体积约六十KB，目录结构清晰，便于快速查阅。已有四百一十四人浏览学习，适合熟悉Python且希望上手频域建模的中高级开发者。项目内附一维至三维FNO变体、低秩分解、超分辨率和图像上采样等模块，并提供了Adam优化器、数据生成脚本以及Navier-Stokes、Burgers、Darcy等经典PDE的求解与评估流程，能够帮助读者从代码层面完整理解傅里叶层设计、快速傅里叶变换、损失函数与网络泛化等核心细节，是开展物理深度学习研究的实用入门项目。

1. 傅里叶神经算子：一个把 PDE 求解器换掉的开源项目包

A同学前两天把一份fourier_neural_operator-master.zip发给我，说老板让他在 128×128 网格上算二维流体演化，传统谱方法每步求解压力泊松方程要几十秒，根本迭代不动。我告诉他，换傅里叶神经算子（Fourier Neural Operator，简称 FNO）这条路：不逐点求解，而是把输入场和输出场都变换到频域，直接在傅里叶模态上做可学习的加权，把整个 PDE 的演化过程当成一个函数学习问题。这份压缩包解压后模型部分其实很短，真正容易卡住的反而是数据格式、维度约定和模态截断这几个隐蔽位置。这篇文章我按实际复现的顺序写，把代码结构、训练流程、调参经验依次展开，手头有类似流体或扩散仿真需求的，可以照着走一遍。

2. FNO 原理与代码结构：频谱卷积从矩阵到 FFT

2.1 从卷积定理到谱卷积：为什么在频域做权重乘法

先放下代码，把 FNO 的核心动机讲清楚。时域卷积等价于频域乘法，这是信号处理里的经典结论：给定核函数，卷积操作在傅里叶变换后会退化成对应频率上的逐点乘积。FNO 把这个结论反着用——不先在空间域定义卷积核，而是直接在傅里叶模态上放置一组可学习的复数权重，让网络自己去学每个波数上应该乘多少。对 PDE 算子学习来说，这个做法的意义很大：流场里的大尺度结构（旋涡、边界层、压力波）天然是全局耦合的，传统 CNN 卷积核只有局部感受野，想传递长距离信息必须堆很深；而谱卷积一次变换就能让每个输出位置看到整场信息，一层就够。

模态的数量对应“截断到多少个频率”。二维场做 rfft2 后得到的是从直流到奈奎斯特频率的复数系数，FNO 只保留前 modes 个低频分量，丢弃高频尾巴。低频分量携带能量主体，决定了场的整体形态；高频分量是细节和噪声。丢弃高频等于给网络一个隐式的平滑正则，这也是 FNO 在流体数据上很稳的原因之一。参数规模只和模态数与通道数有关，和网格点数没有关系——模型不依赖训练时的网格密度，这是 FNO 分辨率不变性的来源。

CNN 和 FNO 的差异值得放一张对比表：

所以 FNO 的调参重心不在卷积核大小和步长，而在模态数、通道数和哪些频率分量真正承载了解的信息。后面第 5 章展开说这组参数的敏感度。

2.2 模型内部分层：提升层、傅里叶层与投影层

标准的 FNO 结构是三段式。提升层把输入物理量（坐标、时间、初始场）从低维映射到 width 维，通常就是一个全连接层加激活函数；中间串若干傅里叶层进行频域加权；最后投影层把 width 维压回输出场的物理量维。真正的算力几乎都花在傅里叶层里。

下面这段是我按这类包的通用逻辑重写的二维谱卷积层，注释里标出了最容易错的地方：

import torch import torch.nn as nn class SpectralConv2d(nn.Module): def __init__(self, in_channels, out_channels, modes1, modes2): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.modes1 = modes1 # 沿第一维保留的模态数 self.modes2 = modes2 # 沿第二维保留的模态数 # 复数权重，形状：(输入通道, 输出通道, 模态1, 模态2) self.scale = 1.0 / (in_channels * out_channels) self.weights1 = nn.Parameter( self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat) ) self.weights2 = nn.Parameter( self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat) ) def forward(self, x): # x 形状：(batch, in_channels, Nx, Ny) batch = x.shape[0] # rfft2 只保留 Ny 维的一半频谱，节省计算量 x_ft = torch.fft.rfft2(x, dim=(-2, -1)) # x_ft 形状：(batch, in_channels, Nx, Ny//2+1) out_ft = torch.zeros( batch, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat, device=x.device ) # 低频块用 weights1 相乘 out_ft[:, :, :self.modes1, :self.modes2] = torch.einsum( "bixy,ioxy->boxy", x_ft[:, :, :self.modes1, :self.modes2], self.weights1 ) # 对应负频率块用 weights2 相乘 out_ft[:, :, -self.modes1:, :self.modes2] = torch.einsum( "bixy,ioxy->boxy", x_ft[:, :, -self.modes1:, :self.modes2], self.weights2 ) # irfft2 回到空间域，s 参数保证输出尺寸与输入一致 return torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)), dim=(-2, -1))

逻辑说明：rfft2 拉开的是复数谱，乘完复数权重后用 irfft2 回到实数空间域。einsum 里的"bixy,ioxy->boxy"表达的是逐位置的矩阵乘法——每个频点上的输入通道向量乘以一个权重矩阵，得到输出通道向量。两个权重矩阵分别处理正频率和负频率的低频分量，这是由 rfft 频谱的共轭对称性决定的。实际组合 FNO 层时，谱卷积分支的输出会和一个 1×1 卷积的线性残差分支相加，再过激活函数；残差分支补上频域截断丢掉的高频信息，同时提供一条稳定的梯度通道。复数权重初值用scale = 1/(in_channels*out_channels)而不是默认的均匀分布，能明显减少训练初期 loss 波动。

参数说明：modes1和modes2是每个方向保留的模态数，64×64 网格时取 12 左右比较稳；channel 数对应 width，经典实验里 64 是一个“够用”的数字，太小欠拟合，太大收益递减且显存涨得明显。

2.3 代码包文件划分：训练入口与工具函数

压缩包解压后，第一件事不是急着跑，而是先看顶层文件清单。这类包的结构相对固定，我列一下典型划分：

第一次接触这种包的人最常犯的错是上来就跑训练脚本，结果报缺数据。正确顺序是先确认数据目录里有没有东西，没有就找工具模块里的下载入口；数据放好后，再打开训练入口确认数据路径常量写的是相对路径还是绝对路径。路径对不上，后面所有报错都会让你误以为是模型写错了。

运行命令也很简单：

# 解压并进入目录 unzip fourier_neural_operator-master.zip cd fourier_neural_operator-master # 二维训练入口，常见命名是 train_2d.py 或 fno_2d_train.py python train_2d.py --batch-size 20 --epochs 500

如果入口脚本不支持命令行参数，就直接打开脚本，把顶部的batch_size、epochs、weight_decay这些常量改成自己的值。建议第一次跑不要碰任何模型结构参数，只改数据路径和轮数，先把基准跑出来，再谈调优。

注意：命令行参数名每个版本可能不一样，优先看脚本里 argparse 的定义，或者顶部常量的注释。

3. 数据准备与训练流程：从 NS 方程数据到 loss 收敛

3.1 数据管线：形状约定与归一化

FNO 的数据不是图片，而是数值模拟生成的流场快照。常见做法是用谱方法求解 Navier-Stokes 方程，生成多组不同初始条件或不同物理系数下的（输入场, 输出场）样本对。每个样本本质是一段网格上的物理量分布，数据落地是二进制数组文件，读进来以后需要按固定 shape 重新组织。

我第一次跑的时候在这里栽过跟头：直接np.load一个 npy 文件得到三维数组，以为是(样本数, 高度, 宽度)，实际文件的最后一个维度是物理量通道数。正确的打开方式是：

import numpy as np import torch def load_ns_data(data_path, num_samples, nx, ny, num_vars): # 数据是纯二进制数组，先用 fromfile 读原始字节 raw = np.fromfile(data_path, dtype=np.float32) # 按固定顺序重排：(样本数, 网格x, 网格y, 物理量数) raw = raw.reshape(num_samples, nx, ny, num_vars) # 输入：除最后一个物理量外的字段；输出：最后一个物理量 x = raw[..., :-1] # 例如 (速度u, 速度v) y = raw[..., -1:] # 例如 涡量场 # FNO 约定：空间维在后，通道维在倒数第二 x = torch.from_numpy(x).permute(0, 3, 1, 2).float() y = torch.from_numpy(y).permute(0, 3, 1, 2).float() return x, y

逻辑说明：二进制文件用np.fromfile按固定 dtype 读进来后，顺序是 line 式的，必须先 reshape 成(样本数, 网格x, 网格y, 物理量数)才能理解。数据集的发布者通常会在说明文件里给出这四个维度的具体数值，num_vars一般是 3（两个速度分量加一个目标场）。

归一化是 FNO 训练绕不开的一步。直接在原始物理量上算 loss，速度场和涡量场的量级差好几个数量级，loss 会被大数值那个物理量主导。常见做法是按通道统计训练集的均值和标准差，输入输出各自标准化：

def normalize(train_x, train_y, val_x, val_y): # 按通道统计，保持空间维结构 mean_x = train_x.mean(dim=(0, 2, 3), keepdim=True) std_x = train_x.std(dim=(0, 2, 3), keepdim=True) mean_y = train_y.mean(dim=(0, 2, 3), keepdim=True) std_y = train_y.std(dim=(0, 2, 3), keepdim=True) train_x = (train_x - mean_x) / std_x train_y = (train_y - mean_y) / std_y val_x = (val_x - mean_x) / std_x val_y = (val_y - mean_y) / std_y return train_x, train_y, val_x, val_y

逻辑说明：统计量只从训练集计算，验证集和测试集沿用同一组均值和标准差，否则验证误差会失真。keepdim 的作用是保留通道维度，后面广播到完整张量时形状才匹配。

3.2 训练循环与验证指标：相对 L2 误差的计算位置

FNO 训练常见的组合是：训练阶段用 MSE 作为 loss，验证阶段用相对 L2 误差度量精度。相对 L2 误差的定义是对每个样本计算预测与目标之间的二范数差，再除以目标自身的二范数，最后在所有样本上取平均。它不受场的绝对幅值影响，比绝对误差更能反映“像不像”。

def rel_l2(pred, true): # pred 和 true 形状相同：(batch, channels, Nx, Ny) diff_norm = (pred - true).pow(2).sum(dim=(-2, -1)).sqrt() true_norm = true.pow(2).sum(dim=(-2, -1)).sqrt() return (diff_norm / true_norm).mean().item()

逻辑说明：sum(dim=(-2, -1))对每个样本的空间维求和，得到向量后逐元素相除，再取平均。验证阶段每次迭代都打印这个值，比盯着训练 MSE 直观得多。

训练循环里值得额外提的是混合精度。二维 64×64 网格不用混精度也能跑，但三维 64×64×64 网格时显存压力很大，开启 AMP 能省一半显存且几乎不掉精度。标准写法：

scaler = torch.cuda.amp.GradScaler() for epoch in range(epochs): model.train() for xb, yb in train_loader: xb, yb = xb.cuda(), yb.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): pred = model(xb, grid) loss = torch.nn.functional.mse_loss(pred, yb) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

逻辑说明：autocast 让 forward 在低精度下执行矩阵运算，GradScaler 负责反向传播时梯度缩放的自动管理。需要注意的坑是：如果 loss 出现 NaN，scaler 会自动跳过当前的 optimizer 步并衰减缩放因子，这时候不要误以为训练恢复正常，要去查数据源。

3.3 训练参数速查：batch_size、学习率与调度器

第一次跑通阶段，参数设置越保守越好。这是我在多个类似项目上验证过的一组稳妥起步值：

学习率这块要特别说一句：FNO 的复数权重对学习率敏感度比普通 CNN 高。1e-3 下 2D 任务通常没问题，3D 任务如果发现 loss 在几百步后开始震荡，先降到 5e-4，不要直接上学习率预热那一套，复杂度不划算。调度器用 StepLR 就够了，ReduceLROnPlateau 在这种周期性数据上反而容易过早降速。

注意：amp 开启时，调度器要在scaler.step之后再调scheduler.step，而且只在 optimizer 真正执行了更新时调用。排查 loss 曲线突变时，先看是不是调度器在某个 step 触发了学习率断崖。

4. 踩坑排查：FNO 训练中最容易翻车的五个点

4.1 数据文件缺失，报 FileNotFoundError

现象：运行训练入口后立即报错，指向数据目录下某个文件不存在，日志里只有路径没有具体提示。

原因：压缩包通常只带代码不带数据，数据需要单独获取；或者训练脚本里写的是相对路径，你从别的目录启动脚本导致路径错位。

解决：先检查工具模块里有没有数据下载入口，有就先运行它；没有就手动下载数据文件，放到训练入口所在目录的 data/ 子目录下，再比对脚本顶部路径常量。路径统一改成绝对路径是最省心的，省得每次切换目录都提心吊胆。我一般会在脚本顶部加一行print(data_path)，让运行日志直接打印实际读取路径，排查路径问题快得多。

4.2 CUDA out of memory，batch 降到 4 还是崩

现象：显存报错出现在 forward 阶段，莫名其妙，模型参数量看着很小。

原因：参数少不代表显存少，FFT 的中间张量在复数域上是同尺寸张量的两倍存储，且傅里叶层的每个频点都要保存中间结果用于反向传播。width=64、batch=20 的二维任务，傅里叶层临时张量轻松吃掉几 GB 显存。

解决：降显存有三个梯度，按顺序执行：先降 batch 到 8；不够就把 width 从 64 降到 32；还不行就开 AMP。三个梯度都走完仍崩，检查是否在无关位置保存了多余张量——比如把整批数据移到了 GPU 上而不是用 DataLoader 逐批加载。

4.3 训练中段 loss 变成 NaN，且不再恢复

现象：前几十步正常，loss 稳步下降，某一步突然变成 NaN，后续所有迭代都是 NaN。

原因：最常见是学习率偏大导致复数权重更新越界；其次是数据里有 inf 或 NaN 没有被标准化过程清洗掉；还有一种隐蔽情况是 modes 参数接近或超过网格奈奎斯特频率，rfft 之后的索引越界产生非法值。

解决：先打印标准化后的输入输出 min 和 max，确认没有脏数据；再核对 modes1/modes2 必须小于等于分辨率的一半，64×64 网格下最大 32，实际取 12 左右最稳；最后把学习率从 1e-3 降到 5e-4，重新跑。复数权重的初始化尺度如果之前改过，恢复到1/(in_channels*out_channels)的默认值再试。

4.4 训练不报错，但 loss 完全不下降

现象：loss 一直趴在高位震荡，怎么看都不像在收敛。

原因：维度约定错位。FNO 的卷积是在空间维上做 FFT，如果数据排列成(batch, Nx, Ny, C)直接喂给网络，模型会拿通道维当空间维做 FFT，语义完全错乱，训练自然不会收敛。

解决：统一按照“空间维在后、通道维在倒数第二”的约定组织输入。具体做法是在数据加载函数的 return 前打印一次x.shape，确认形状和模型 forward 里期望的形状一致。形态对不上的话，用permute调整而不是view，view 会打乱数据排列顺序导致静默错误。

4.5 低分辨率收敛很好，高分辨率推理误差飙升

现象：64×64 训练集上相对 L2 误差在 1% 附近，直接拿 128×128 的输入跑推理，误差涨到 10% 以上。

原因：FNO 的分辨率不变性是结构性质，模型能接收不同分辨率输入不代表它能在未见过的频段上表现好。训练数据只覆盖了 64×64 网格对应的频谱窗口，高分辨率输入携带的高波数能量是训练时从未见过的分布外数据。

解决：不要指望纯外推白拿高精度。想在高分辨率上部署，要么把一部分训练样本直接在高分辨率网格下生成并混入训练集，要么做数据增强，在训练时对场做随机裁剪和扰动，强迫模型对不同频段都鲁棒。先用一个高分辨率小测试集量一下差距，再决定是否需要回炉。

5. 参数与收敛调优：模态数、通道数和激活函数怎么试

5.1 模态数的物理含义与扫描方法

模态数直接决定模型保留多少频率分量。取少了，输出场会过度光滑，小尺度涡细节丢失，误差主要来自高频；取多了，模型开始拟合训练集里的数值噪声，验证误差不降反升，甚至引入谱混叠。

一个直观经验：64×64 网格上，modes 从 8 到 16 之间是甜点区，12 是一个很稳的中点。做扫描实验时，固定 width 和其他条件，只改 modes，跑完记录每个配置的验证相对 L2：

for modes in [8, 12, 16]: model = FNO2d(modes1=modes, modes2=modes, width=64) val_err = train_and_evaluate(model) print(f"modes={modes}, val_rel_l2={val_err:.4f}")

逻辑说明：这个循环替你做了消融实验。观察点是：误差从 8 提到 12 时通常有明显下降，从 12 提到 16 时可能只有微小提升，如果提升后误差反而上升，说明已经过拟合高频噪声，退回 12。更高分辨率网格下可以按比例放大，但不要超过奈奎斯特频率的一半。

5.2 通道数上界与收益递减

width 决定傅里叶层里每个频点上矩阵乘法的宽度。width 太小，模型表达能力不够，误差压不下去；width 太大，参数量和显存都涨，但精度收益不是线性的。

我复现时的印象是：width 从 20 提到 64，验证误差下降很明显；从 64 提到 128，误差下降幅度小一个数量级，显存占用却近乎翻倍。二维任务里 width=64 是性能与资源的最佳平衡点，三维任务由于显存压力大，width=32 更实际。

表格说明：这不是精确的基准数据，只是给你一个相对量级的参考。调参时先固定 width=64 跑通全流程，后面再看资源余量决定要不要加宽。

5.3 激活函数与训练配置的交互

FNO 层内激活函数的选择不如 CNN 里那么敏感，但也不是可以随手乱换。ReLU 在频域边界上会引入轻微的锯齿效应，但收敛速度快；GELU 更平滑，训练曲线更稳定，代价是每个傅里叶层多一点点计算量。我的习惯是：2D 快速实验用 ReLU，3D 或高分辨率任务用 GELU，图个稳定。

优化器方面，Adam 配weight_decay取 1e-4 或干脆关闭，FNO 的复数权重本身自带正则效果，加 L2 反而容易把低频分量压过头。梯度裁剪默认不开，真的发散时先怀疑学习率，不要用 clip 掩盖问题。调度器的使用节奏前面提过，StepLR 按固定步数减半就够，关键是减半时机要落在 loss 平台期之后，不然会把模型推到过小的学习率区间里动弹不得。

6. 一个验证技巧：用分辨率外推检验模型真的学到了算子

训练结束、误差也好看，这时候最该做的不是庆祝，而是问一个问题：模型是学到了算子映射，还是记住了训练集里的样本？分辨率的零样本外推测试是回答这个问题最直接的办法。训练时用 64×64，推理时把输入场在更细网格上重新计算或插值到 128×128，喂给同一个模型：

model.load_state_dict(torch.load("checkpoint_best.pt", map_location="cpu")["model"]) model = model.cuda().eval() with torch.no_grad(): pred_hr = model(x_hr.cuda(), grid_hr.cuda()) err_hr = rel_l2(pred_hr.cpu(), y_hr) print(f"128x128 外推相对误差: {err_hr:.4f}")

判别准则：如果外推误差与训练分辨率下的误差同量级，说明模型学到的是输入到输出的函数映射，频域结构理解到位；如果误差翻了好几倍，说明模型大概率依赖了网格密度相关的特征，FNO 的结构优势没有真正发挥。第二种情况不代表模型废了，而是训练数据的频谱覆盖不够，回炉加高分辨率样本即可。

再配合一个频谱检查，判断模型是在学物理还是在学图像纹理：把预测场和真值场分别做 rfft2，按半径分桶统计能量，比较径向谱曲线。低频段重合良好、高频段偏离是正常形态；如果低频段就劈叉，说明模态截断设置过小，连主体能量都没保住。这两步做完，模型的可靠性心里基本有数了。

从那以后我拿到任何一份训练好的 FNO 权重，第一件事不是看 loss 曲线，而是把这个分辨率外推测试跑一遍，再决定信不信它。希望帮到你。

本文还有配套的精品资源，点击获取