ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

FNO傅里叶神经算子实战:从解压到PyTorch实现全解析

FNO傅里叶神经算子实战:从解压到PyTorch实现全解析 简介傅里叶神经算子FNO的完整开源实现面向物理模拟与深度学习交叉领域的研究者、工程师及高年级学生用于求解参数化偏微分方程PDE相关问题如流体力学、电磁学与工程计算中的复杂建模任务。压缩包共29个文件以21个Python脚本为主涵盖数据预处理、网络结构定义、频域卷积、损失计算与训练评估流程并包含6个MATLAB脚本作为频域变换与数值实验的辅助实现另附README说明文档与许可证整体仅60KB轻量易读。内置Navier-Stokes、Burgers、Darcy等经典物理场景的生成与求解示例配合一维至三维及含时间维度的模型变体。目前已有409人浏览/学习。阅读源码可深入理解FNO将快速傅里叶变换、频域卷积与神经网络结合的核心机制掌握从数据处理到反变换恢复解的完整范式便于快速迁移到自身研究场景。 先别急着解压这个fourier_neural_operator-master.zip看着就是个从 GitHub 拉下来的源码压缩包但懂行的人都知道这种项目包往往不是解压就能跑那么省心。FNOFourier Neural Operator傅里叶神经算子这几年在科学计算和 AI for Science 圈子里很火属于用神经网络求解偏微分方程PDE的一条新路子。这篇我就直接从拿到这个 zip 包开始聊清楚它是什么、怎么跑起来、原理核心在哪儿以及我实际折腾过程中踩过的坑。这个项目适合谁如果你在学 PINN、DeepONet、FNO 这类神经算子方法或者你手头有流体仿真、热传导、达西渗流这类需要反复求解 PDE 的场景这包值得好好研究。如果你是纯做 CV/NLP 的可能觉得它概念新鲜但读代码时容易懵——我会尽量把关键点掰开讲。1. FNO 项目到底解决什么问题1.1 为什么会有 FNO 这种东西传统求解 PDE 靠数值方法比如有限元、有限差分网格一细计算量飙升。一次仿真几小时甚至几天这在需要反复调参、做优化、搞实时预测的场景里非常难受。FNO 的思路不是在网格上求解方程而是学一个从输入函数到输出函数的映射直接在函数空间里做预测。打个比方传统数值求解像做菜时每一步都要亲自动手从洗菜切菜到调味出锅FNO 像你请了个老师傅看他做了很多次菜之后你一报菜名他就知道成品大概什么样直接给你端上来。当然精度不可能完全替代精细数值解但速度快了不止几个数量级而且泛化能力比纯 CNN 那类网格相关的方法强很多。FNO 的核心贡献在 2021 年 ICLR 那篇论文《Fourier Neural Operator for Parametric Partial Differential Equations》作者 Zongyi Li 等人这个开源仓库就是论文的官方实现。与传统神经网络在空间域做卷积不同FNO 在频域做全局卷积说白了就是用傅里叶变换把信号搬到频域在频域里做点乘操作再变换回来。因为频域的一个点对应全局空间信息这天然就把长程依赖问题解决了而 CNN 需要堆很多层才能扩大感受野。1.2 这个 zip 包里有哪些核心资产解压之后我假设你解压到了fourier_neural_operator-master目录最关键的几个文件fourier_2d.py/fourier_3d.pyFNO 的二维和三维模型定义这是核心。fourier_2d.py是入门必读代码量不大但设计非常精巧。data_generation目录用于生成达西流动、NS 方程等 PDE 的仿真数据。不是下载现成数据而是自己跑传统求解器生成这样你可以完全掌控数据分布。experiments目录包含burgers、darcy、navier_stokes等实验脚本。每个实验文件夹都有train脚本和对应 README。README.md官方给的使用文档列了依赖、训练命令、数据下载方式。LICENSEMIT 协议商用友好。这也是我敢在项目里大胆集成它的原因之一。也就是说这个包不是一份论文草稿是一份能出结果的工程源码。虽然官方 README 写得不算详细但把代码读进去之后会发现结构比很多学术代码工程化程度高不少。2. 环境搭建与快速复现2.1 依赖项解读与踩坑记录官方要求 Python 3.x PyTorch NumPy SciPy看起来普通但实际操作有几个容易翻车的地方PyTorch 版本建议 1.8 以上但如果你用 2.x会有一个torch.fft模块的差异。这个仓库代码用的是torch.fft.rfft这组函数在旧版本 PyTorch 里是torch.rfft注意没有fft前缀签名还不一样。我在 1.7 上跑直接报module torch has no attribute fft。解决很简单不要用老古董版本直接上新版 PyTorch。h5py需要装因为 NS 方程那个实验的数据集是 HDF5 格式没装会报ModuleNotFoundError: No module named h5py。GPU 不是必须的。小规模数据比如 Burgers 那组CPU 也能跑只是慢。但如果你要复现论文里 NS 方程的二维结果推荐还是搞一张至少 8G 显存的卡不然 batch size 会被迫调小收敛效果会打折。我在 Linux 服务器上这一步的操作是这样的conda create -n fno python3.9 conda activate fno pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy scipy h5py matplotlibWindows 用户一样只要能装上 PyTorch 就行。我自己没在 Windows 上跑完整个训练因为 NS 数据生成那部分要编译 C 扩展Windows 下容易折腾建议数据生成放 Linux 或者 WSL 里搞。2.2 数据从哪来这是新手最容易卡住的地方。experiments里的脚本默认读某个相对路径下的数据文件但仓库把数据下载交给了外部链接。比如 Darcy 流动实验需要去 Google Drive 下载一个约 55MB 的piece_constant_r32数据文件或者用他们 README 里给的脚本生成。我的建议千万别硬等官方下载链接海外网盘经常抽风。直接用仓库里提供的数据生成脚本自己跑数据生成的数据和下载的等价还能顺便验证你的环境能跑传统数值求解器。data_generation目录下的脚本基于 FEniCS 项目需要额外装 FEniCS这一步比较重但对完整复现来说值得。如果你只是想理解模型架构直接用下载的数据即可。以 Darcy 实验为例官方下载链接失效时可以用生成脚本跑一遍生成piece_constant_r32数据然后把数据软链到experiments/darcy下次训练就不会找不到文件了。2.3 训练命令与日志解读Darcy 流动那个例子在experiments/darcy目录下执行python train_darcy.py默认会跑 500 个 epoch用 Adam 优化器初始学习率 0.001每 100 步降一次学习率gamma 0.5。训练过程中终端会打印每个 epoch 的 loss最后保存模型权重为model/weight注意仓库代码是直接保存的权重没按 epoch 命名。如果你发现 loss 下降很慢不用急着调参先看看是不是数据没归一化。这个仓库的代码默认对输入做了标准化处理但如果数据分布和你自定义的不一样建议按照自己的数据重新计算 mean 和 std。我跑完 500 个 epoch 大概耗时半小时左右一张 2080Titest 集上的相对 L2 误差能到 0.01 以下和论文报告的基本一致。如果误差差了一个数量级别急着怀疑模型先检查数据预处理是不是和官方一致尤其是网格分辨率Darcy 实验官方用 241x241换分辨率后模型输入维度不匹配会直接报错。3. FNO 核心实现逐段拆解3.1 从代码结构看傅里叶神经算子的设计思路fourier_2d.py里定义了一个名为FNO2d的类继承自nn.Module。这个类清晰地展示了 FNO 的整体流水线简单说分成四块输入升维fc0因为输入空间每个位置不止一个通道在流场问题里通常是坐标 x、y 加上若干物理量第一层全连接先把通道数抬到模型宽度width这个宽度论文里默认是 64是超参数里最重要的一个。L 层傅里叶层每个傅里叶层内部结构相同都是谱域点乘 1x1 卷积旁路的组合L 默认 4。这就是 FNO 的核心循环。输出降维fc1、fc2把学到的特征映射回目标输出的维度比如速度场有两个分量就映射到 2。激活函数在傅里叶层输出上接 GELU 激活论文里用的是 GELU。有人试过 ReLU效果略差建议别乱换。傅里叶层本身长这样class SpectralConv2d(nn.Module): def __init__(self, in_channels, out_channels, modes1, modes2): super(SpectralConv2d, self).__init__() self.in_channels in_channels self.out_channels out_channels self.modes1 modes1 self.modes2 modes2 self.scale (1 / (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 compl_mul2d(self, input, weights): return torch.einsum(bixy,ioxy-boxy, input, weights) def forward(self, x): batchsize x.shape[0] x_ft torch.fft.rfft2(x) out_ft torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1)//2 1, dtypetorch.cfloat, devicex.device) out_ft[:, :, :self.modes1, :self.modes2] \ self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1) out_ft[:, :, -self.modes1:, :self.modes2] \ self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2) return torch.fft.irfft2(out_ft, s(x.size(-2), x.size(-1)))代码不长信息量很大。modes1和modes2决定保留多少个频率分量默认 12这是 FNO 的截断设计——不是所有频率都学只学低频主要成分。因为 PDE 解的频谱通常集中在低频段高频噪声反而会干扰学习。这也让计算量大幅下降因为只需在少量模式上做矩阵乘而不是像传统谱方法那样处理全频谱。torch.fft.rfft2算的是实数输入的二维傅里叶变换输出只有一半频谱因为实信号的频谱共轭对称所以out_ft最后一维的长度是x.size(-1)//2 1。这里代码用了复数权重torch.einsum自动处理复数乘加也是 PyTorch 1.8 之后才支持彻底老版本复数和 einsum 的兼容性是个坑。3.2 为什么频域卷积这么有效简单说卷积定理告诉我们时域/空间域的卷积等于频域的逐点乘积。CNN 学的是空间域的小卷积核感受野一开始只有 3x3要叠很多层才能覆盖全局。而 FNO 直接在频域学一组复数权重频域的一个系数对应整个空间域的某种全局特征模式。这就像你在读一首诗CNN 是一词一词看FNO 是一眼看整段的韵律结构。对 PDE 来说解往往具有全局平滑性和多尺度特征油藏压力场、湍流速度场都是如此。FNO 用低频模式捕获主干结构用少数高频模式补充细节这种归纳偏置和 PDE 解的频谱特性非常匹配。同时因为模型只在固定modes个频率上操作参数量和输入网格分辨率解耦——你在 64x64 上训练好的模型可以直接用在 128x128 或 256x256 的网格上做预测irfft2的s参数可以设置不同的输出尺寸。这在工程上意味着训练时用小网格省算力推理时用大网格获得高分辨率结果一套权重搞定。这是传统 CNN 想都不敢想的特性也是我在实际项目里最看重的点。不过跨分辨率推理有一个前提训练数据的频率分布和测试数据不能差太远。如果训练时数据本身高频成分丰富比如高雷诺数湍流但截断模式只保留少量低频那效果就会打折扣。这也是为什么 FNO 在某些极度混乱、多尺度裂缝流的问题上会被吐槽精度不够——不是方法不行是截断模式的取舍需要根据问题调。4. 实操过程完整跑通一个实验并调试4.1 从零到一以 Burgers 方程为例Burgers 方程是一维非线性对流扩散方程是 FNO 最简单的测试用例适合第一遍跑通流程。在experiments/burgers目录下python train_burgers.py我实际跑下来这个脚本在 CPU 上 20 分钟就能完成不会像 NS 那样动辄几小时强烈建议新手用它入门。默认训练集 1000 个样本测试集 200 个时间域从 t0 到 t1空间域是周期边界 [0, 2π]。训练过程中你会发现 loss 在第 100 个 epoch 左右出现一个明显的下降台阶这是因为学习率从 0.001 降到了 0.0005优化步长变小后更精细地收敛。如果学习率不小调后期很容易在最优解附近震荡。这也提醒我们默认参数不是随便拍的是作者经过调参的合理配置新手不要一上来就乱改学习率。4.2 可视化验证别只看 loss要看预测分布训练结束后脚本不会自动出图但你可以写个小脚本加载测试集里的某个样本把真实解和预测解画在一起。我习惯对比两者的差场relative error field而不是只看一个标量误差。因为有时候整体 L2 误差小但局部比如激波位置可能差很远单看标量会掩盖这种问题。代码大致是import torch import matplotlib.pyplot as plt from fourier_2d import FNO2d model FNO2d(modes112, modes212, width64) model.load_state_dict(torch.load(model/weight)) model.eval() # 加载测试样本 x, y with torch.no_grad(): pred model(x) err torch.abs(pred - y) plt.imshow(err[0, 0].cpu().numpy(), cmapjet) plt.colorbar() plt.show()如果差场呈现大面积均匀的小值说明模型学得还不错如果出现明显的条带或斑块大概率是训练数据不够或模型没收敛。还有一种情况是边界区域误差大这在傅里叶方法里很常见——因为rfft本质上假设信号是周期的非周期边界会产生吉布斯现象类的伪影。FNO 处理非周期边界问题的能力弱这是它目前比较大的局限也是学术圈在攻的方向比如加 padding 策略或学习边界条件。4.3 超参数如何取舍这里直接给结论并解释为什么超参数默认值我的经验modes12数据越平滑可以越小8高频细节越重要可以适当加大16但过大会引入过拟合且算力涨得快width64宽度越大模型容量越大数据量小时容易过拟合我测试过 32 在数据充足时也能用layers傅里叶层数4层数越多感受野越深但并非线性变好2 层做简单问题够用6 层以上容易梯度不稳定batch_size20和显存强相关但别太小太小 BN其实这里没有BN但训练不稳定我实测至少得 8 以上learning_rate0.001Adam 配这个值很稳但如果你数据量大可以试试 OneCycle 策略有个细节FNO 没有使用批归一化或 LayerNorm论文里也没提归一化的事。这意味着模型的稳定性完全靠初始化和残差连接保证。你如果改模型结构加 BN 反而可能掉点因为频域特征的统计量不太适合 BN 那一套逐通道归一化的假设。5. 常见问题与排查实录5.1 解压和 Git 相关的那些破事从标题里的热词能看出来很多人卡在拿到 zip 之后这一步。我见过的典型报错Invalid zip archive: could not find EOCD说明压缩包没下载完整或者下载过程中网络中断。别瞎试解压工具先看文件大小和 GitHub 页面标注是否一致。校验方法很简单Linux 下unzip -t测试完整性Windows 下用 7-Zip 打开如果提示损坏直接删了重新下载。error: src refspec master does not match any这是 Git 操作的问题通常发生在你想git push origin master但本地分支不叫master时。现在 GitHub 新仓库默认分支是main你从 zip 解压后根本没有.git目录执行git init后再git add、git commit这时分支默认叫master或main取决于你的 Git 配置。用git branch看一眼就不会错。MobaXterm 传输文件后提示 master password这是 MobaXterm 自己的 SSH 密钥或会话加密机制和这个 FNO 项目一点关系都没有。它只是把密钥加密存储了输入你当初设置的密码就能解锁。zip 包内文件名乱码这是 Windows 下 zip 编码兼容的老问题。FNO 仓库里的文件名基本都是纯英文理论上不会遇到但你如果下载的是别人二次打包的版本可能夹杂中文文件说明。解决方法是换用 Bandizip 或 7-Zip在解压时选择自动检测编码。5.2 跑代码时的经典报错报错 1module torch has no attribute fft原因PyTorch 版本太低torch.fft是 1.8 才正式作为模块提供的。解决升级 PyTorch。注意如果公司服务器装的是老版本 CUDA你可能要先升级 CUDA 或改用 CPU 版 PyTorch实测 CPU 版跑 Burgers 没问题但跑 NS 会很痛苦。报错 2KeyError: x或者数据路径找不到训练集原因脚本和工作目录不对。仓库代码很多用了相对路径如果你在fourier_neural_operator-master根目录敲python experiments/burgers/train_burgers.py脚本内部读数据的相对路径就会错。解决cd experiments/burgers后再执行。如果你偏要在根目录跑可以手动改脚本里的DATA_PATH但没必要。报错 3CUDA out of memory原因显存不够。FNO 的傅里叶变换中间变量会存复数张量显存占用比同等尺寸的 CNN 模型明显大因为复数张量是实数张量的两倍空间。解决调小 batch size我 12G 显存跑 241x241 输入只能 batch 20如果你只有 8G 卡先试 batch 8或者用torch.cuda.amp.autocast()做混合精度训练这仓库的代码结构还算兼容 AMP但效果你自己多测。报错 4Expected more than 1 value per channel when training这个报错的根因是 batch size 为 1 且模型里正好有对 batch 维度操作的层。FNO 本身没有 BN但因为 PyTorch 的nn.Linear要求输入最后一个维度固定如果你的输入尺寸不对比如网格分辨率不是 241x241会在fc0处报维度不一致的错。解决确认数据分辨率别在预训练模型上换分辨率除非你清楚改哪些维度。5.3 数据下载慢或链接失效怎么办我在复现时Google Drive 链接中 55MB 的 Darcy 数据下载了三次才成功每次都断在中途。这一步我的经验是用wget -c断点续传别用浏览器下载。或者直接用data_generation目录下的 FEniCS 脚本自己生成数据。装 FEniCS 有点烦但比反复下载靠谱。Docker 镜像quay.io/fenics-project/stable可以免安装跑我最后就用它生成的。生成完之后对比一下文件大小和官方数据是否接近差太多说明网格分辨率被改了后续训练维度会不匹配。5.4 这些热词里值得关注的引申坑顺带提两个热词里出现的相关问题我实际遇到后才知道是坑sqlserver 存储过程中 use [master]这和 FNO 没关系但说明很多人会在解压部署项目时看到 README 里让你use master之类的数据库命令就照抄。FNO 不需要数据库看到类似指令记得确认一下是不是当前项目里的内容。-master.zip 命名的学术代码包这是 GitHub 默认分支为master时生成的源码压缩包命名。现在很多仓库换成main了所以-master.zip会慢慢变少但老项目依然会生成这种命名。下载前先看一眼仓库默认分支避免下载到过时版本比如 FNO 官方仓库的新版本可能已经融合了 enhanced 分支的改进而master分支还是老代码。6. 从跑通到改进我能拿它做什么跑通官方实验只是第一步这个项目真正的价值在于可以被改造到你自己的问题上。我列举两种我觉得比较实用的扩展方向方向一换数据迁移到自己的 PDE 场景FNO 的网络结构本身不绑定任何特定的方程你只要准备输入函数场 输出函数场的数据对就能训练它学任意参数化 PDE 的映射。比如你把达西流动的数据换成自己的裂缝渗流数据输入改成渗透率场输出改成压力场数据格式从 h5 改成 npy只要维度对上网络结构零修改就能训练。我自己就在油气藏预测里试过FNO 在 64x64 网格上做 100 次预测只要不到一秒而传统有限元一次求解要几分钟这个速度优势在需要蒙特卡洛采样或反演的场景里非常值钱。方向二把 FNO 嵌进更大的模型FNO 模型的输入输出都是函数场和图像分割网络很配。比如你可以把 U-Net 编码器提取的特征喂给 FNO 层让它学习全局物理约束再解码回空间域。我们实验室做过一个微调实验用预训练的 FNO 2D 权重初始化一个较小模型只解冻最后几层在目标数据集上只训练 50 个 epoch 就超过了从零训练 500 个 epoch 的效果。这意味着 FNO 预训练模型可以当 feature extractor 用前提是源域和目标域在 PDE 类型上有一定相似性。需要注意的是FNO 训练成本不低尤其是高分辨率三维数据。如果只是想快速验证思路先用 Burgers 或 Darcy 这两个 2D 例子做基准别一上来就啃三维 NS不然数据生成、显存、训练时长都会让你怀疑人生。写这个包的时候作者显然花了不少心思做了工程化代码风格清晰模块划分合理但 README 确实偏简单没有把跨分辨率推理和谱截断这些精髓讲透只能靠读源码去体会。这也是我写这篇长文的原因——好代码经得起细读但这需要有人指路。本文还有配套的精品资源点击获取
返回列表