简介傅里叶神经算子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 原理与代码结构频谱卷积从矩阵到 FFT2.1 从卷积定理到谱卷积为什么在频域做权重乘法先放下代码把 FNO 的核心动机讲清楚。时域卷积等价于频域乘法这是信号处理里的经典结论给定核函数卷积操作在傅里叶变换后会退化成对应频率上的逐点乘积。FNO 把这个结论反着用——不先在空间域定义卷积核而是直接在傅里叶模态上放置一组可学习的复数权重让网络自己去学每个波数上应该乘多少。对 PDE 算子学习来说这个做法的意义很大流场里的大尺度结构旋涡、边界层、压力波天然是全局耦合的传统 CNN 卷积核只有局部感受野想传递长距离信息必须堆很深而谱卷积一次变换就能让每个输出位置看到整场信息一层就够。模态的数量对应“截断到多少个频率”。二维场做 rfft2 后得到的是从直流到奈奎斯特频率的复数系数FNO 只保留前 modes 个低频分量丢弃高频尾巴。低频分量携带能量主体决定了场的整体形态高频分量是细节和噪声。丢弃高频等于给网络一个隐式的平滑正则这也是 FNO 在流体数据上很稳的原因之一。参数规模只和模态数与通道数有关和网格点数没有关系——模型不依赖训练时的网格密度这是 FNO 分辨率不变性的来源。CNN 和 FNO 的差异值得放一张对比表对比项CNNFNO感受野局部需深网络叠加全局单层覆盖整场参数与网格关系卷积核尺寸固定与网格无关参数在频域与网格无关分辨率外推特征尺度漂移易失效保持算子映射可外推数值风格偏视觉任务调优偏光谱分析与模态截断所以 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, dtypetorch.cfloat) ) self.weights2 nn.Parameter( self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtypetorch.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//21) out_ft torch.zeros( batch, self.out_channels, x.size(-2), x.size(-1) // 2 1, dtypetorch.cfloat, devicex.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 代码包文件划分训练入口与工具函数压缩包解压后第一件事不是急着跑而是先看顶层文件清单。这类包的结构相对固定我列一下典型划分文件位置常见职责使用时机顶层训练入口二维或三维训练脚本负责数据加载、训练循环、日志输出项目主入口模型定义模块FNO 网络结构、谱卷积层、提升与投影层被训练入口导入工具模块数据读取、标准化、损失计算、网格坐标生成被训练入口调用数据目录存放数值模拟生成的流场数据训练前必须就位检查点目录保存模型权重和优化器状态训练中定期写入第一次接触这种包的人最常犯的错是上来就跑训练脚本结果报缺数据。正确顺序是先确认数据目录里有没有东西没有就找工具模块里的下载入口数据放好后再打开训练入口确认数据路径常量写的是相对路径还是绝对路径。路径对不上后面所有报错都会让你误以为是模型写错了。运行命令也很简单# 解压并进入目录 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, dtypenp.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), keepdimTrue) std_x train_x.std(dim(0, 2, 3), keepdimTrue) mean_y train_y.mean(dim(0, 2, 3), keepdimTrue) std_y train_y.std(dim(0, 2, 3), keepdimTrue) 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 出现 NaNscaler 会自动跳过当前的 optimizer 步并衰减缩放因子这时候不要误以为训练恢复正常要去查数据源。3.3 训练参数速查batch_size、学习率与调度器第一次跑通阶段参数设置越保守越好。这是我在多个类似项目上验证过的一组稳妥起步值参数推荐起步值说明batch_size2D16203D48显存不够先降它学习率1e-3Adam 的默认配合值调度器StepLR(step_size500, gamma0.5)500 步减半适合长训练epochs500FNO 收敛偏慢别设 50width643D 任务降到 32modes1/modes21264×64 网格下比较中庸损失函数MSE验证阶段再看相对 L2学习率这块要特别说一句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 memorybatch 降到 4 还是崩现象显存报错出现在 forward 阶段莫名其妙模型参数量看着很小。原因参数少不代表显存少FFT 的中间张量在复数域上是同尺寸张量的两倍存储且傅里叶层的每个频点都要保存中间结果用于反向传播。width64、batch20 的二维任务傅里叶层临时张量轻松吃掉几 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调整而不是viewview 会打乱数据排列顺序导致静默错误。4.5 低分辨率收敛很好高分辨率推理误差飙升现象64×64 训练集上相对 L2 误差在 1% 附近直接拿 128×128 的输入跑推理误差涨到 10% 以上。原因FNO 的分辨率不变性是结构性质模型能接收不同分辨率输入不代表它能在未见过的频段上表现好。训练数据只覆盖了 64×64 网格对应的频谱窗口高分辨率输入携带的高波数能量是训练时从未见过的分布外数据。解决不要指望纯外推白拿高精度。想在高分辨率上部署要么把一部分训练样本直接在高分辨率网格下生成并混入训练集要么做数据增强在训练时对场做随机裁剪和扰动强迫模型对不同频段都鲁棒。先用一个高分辨率小测试集量一下差距再决定是否需要回炉。5. 参数与收敛调优模态数、通道数和激活函数怎么试5.1 模态数的物理含义与扫描方法模态数直接决定模型保留多少频率分量。取少了输出场会过度光滑小尺度涡细节丢失误差主要来自高频取多了模型开始拟合训练集里的数值噪声验证误差不降反升甚至引入谱混叠。一个直观经验64×64 网格上modes 从 8 到 16 之间是甜点区12 是一个很稳的中点。做扫描实验时固定 width 和其他条件只改 modes跑完记录每个配置的验证相对 L2for modes in [8, 12, 16]: model FNO2d(modes1modes, modes2modes, width64) val_err train_and_evaluate(model) print(fmodes{modes}, val_rel_l2{val_err:.4f})逻辑说明这个循环替你做了消融实验。观察点是误差从 8 提到 12 时通常有明显下降从 12 提到 16 时可能只有微小提升如果提升后误差反而上升说明已经过拟合高频噪声退回 12。更高分辨率网格下可以按比例放大但不要超过奈奎斯特频率的一半。5.2 通道数上界与收益递减width 决定傅里叶层里每个频点上矩阵乘法的宽度。width 太小模型表达能力不够误差压不下去width 太大参数量和显存都涨但精度收益不是线性的。我复现时的印象是width 从 20 提到 64验证误差下降很明显从 64 提到 128误差下降幅度小一个数量级显存占用却近乎翻倍。二维任务里 width64 是性能与资源的最佳平衡点三维任务由于显存压力大width32 更实际。width相对 L2 误差量级显存占用单轮耗时20较高欠拟合明显低快64收敛良好甜点区中等合理128收益递减高明显变慢表格说明这不是精确的基准数据只是给你一个相对量级的参考。调参时先固定 width64 跑通全流程后面再看资源余量决定要不要加宽。5.3 激活函数与训练配置的交互FNO 层内激活函数的选择不如 CNN 里那么敏感但也不是可以随手乱换。ReLU 在频域边界上会引入轻微的锯齿效应但收敛速度快GELU 更平滑训练曲线更稳定代价是每个傅里叶层多一点点计算量。我的习惯是2D 快速实验用 ReLU3D 或高分辨率任务用 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_locationcpu)[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(f128x128 外推相对误差: {err_hr:.4f})判别准则如果外推误差与训练分辨率下的误差同量级说明模型学到的是输入到输出的函数映射频域结构理解到位如果误差翻了好几倍说明模型大概率依赖了网格密度相关的特征FNO 的结构优势没有真正发挥。第二种情况不代表模型废了而是训练数据的频谱覆盖不够回炉加高分辨率样本即可。再配合一个频谱检查判断模型是在学物理还是在学图像纹理把预测场和真值场分别做 rfft2按半径分桶统计能量比较径向谱曲线。低频段重合良好、高频段偏离是正常形态如果低频段就劈叉说明模态截断设置过小连主体能量都没保住。这两步做完模型的可靠性心里基本有数了。从那以后我拿到任何一份训练好的 FNO 权重第一件事不是看 loss 曲线而是把这个分辨率外推测试跑一遍再决定信不信它。希望帮到你。本文还有配套的精品资源点击获取
