Appearance
最快的路径不是先系统学习 PDE,也不是先读几十篇论文,而是:
先把 FNO 当作一个“输入场 → 输出场”的 PyTorch 模型跑通,再逐步拆开其内部结构。
对于你的目标:AI for Science、工程仿真、后续做公开项目,建议采用 复现 → 改写 → 自己实现 → benchmark 四阶段路线。
1. 先明确:你到底在训练什么?
普通神经网络通常学习有限维映射:
Neural Operator 学习的是函数之间的映射:
例如 Darcy Flow:
输入
a(x, y):多孔介质的渗透率场;输出
u(x, y):对应的压力场;每个样本不是一个向量,而是一张二维网格图。
在 PyTorch 中,暂时不必把它想得过于抽象。你处理的张量仍然是:
python
x.shape == [batch_size, in_channels, height, width]
y.shape == [batch_size, out_channels, height, width]区别是:模型试图学习一族 PDE 实例的解算子,而不是拟合某一个固定 PDE 实例的单一解。
2. 第一天:直接跑通官方 FNO Darcy Flow 示例
优先使用 neuraloperator/neuraloperator。这是当前维护中的 PyTorch Neural Operator 库,包含官方 FNO 实现,已经进入 PyTorch Ecosystem。官方文档提供完整的 Darcy Flow 示例,包括数据加载、训练、损失函数、可视化和 zero-shot super-resolution。(GitHub)
2.1 安装
在 WSL2 Ubuntu 或服务器上创建独立环境:
bash
conda create -n neuralop python=3.11 -y
conda activate neuralop
pip install torch torchvision
pip install neuraloperator需要修改库源码时,再改为 editable install:
bash
git clone https://github.com/neuraloperator/neuraloperator.git
cd neuraloperator
pip install -e .
pip install -r requirements.txt这两种安装方式都是官方仓库给出的路径。(GitHub)
2.2 先运行这个最小脚本
创建 run_fno_darcy.py:
python
import torch
from neuralop import H1Loss, LpLoss, Trainer
from neuralop.data.datasets import load_darcy_flow_small
from neuralop.models import FNO
from neuralop.training import AdamW
device = "cuda" if torch.cuda.is_available() else "cpu"
train_loader, test_loaders, data_processor = load_darcy_flow_small(
n_train=1000,
batch_size=64,
n_tests=[100, 50],
test_resolutions=[16, 32],
test_batch_sizes=[32, 32],
)
data_processor = data_processor.to(device)
model = FNO(
n_modes=(8, 8),
in_channels=1,
out_channels=1,
hidden_channels=24,
projection_channel_ratio=2,
).to(device)
optimizer = AdamW(
model.parameters(),
lr=1e-2,
weight_decay=1e-4,
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=30,
)
h1_loss = H1Loss(d=2)
l2_loss = LpLoss(d=2, p=2)
trainer = Trainer(
model=model,
n_epochs=15,
device=device,
data_processor=data_processor,
wandb_log=False,
eval_interval=5,
use_distributed=False,
verbose=True,
)
trainer.train(
train_loader=train_loader,
test_loaders=test_loaders,
optimizer=optimizer,
scheduler=scheduler,
regularizer=False,
training_loss=h1_loss,
eval_losses={
"h1": h1_loss,
"l2": l2_loss,
},
)运行:
bash
python run_fno_darcy.py官方示例使用小型 Darcy Flow 数据集:训练分辨率为 16 × 16,并在 16 × 16 与 32 × 32 上测试。其目的不是达到研究级精度,而是快速展示完整工作流,以及 FNO 的 zero-shot super-resolution:训练时只见过低分辨率数据,推理时仍可直接处理更高分辨率网格。官方文档中的示例在 CPU 上即可快速完成。(Neural Operator)
3. 第一阶段只需要理解 6 个对象
不要一开始钻进全部源码。先理解数据流。
| 对象 | 作用 | 初期需要掌握什么 |
|---|---|---|
train_loader | 提供 PDE 输入输出场 | 检查 x.shape、y.shape |
data_processor | 标准化、预处理 | 知道训练前后做了变换 |
FNO | 学习算子映射 | 理解输入输出 channel 和 Fourier modes |
n_modes | 保留的 Fourier 模态数量 | 相当于频域带宽 |
H1Loss | 比较函数值和梯度 | PDE 场预测常用 |
Trainer | 封装训练循环 | 跑通后应自行拆掉 |
官方示例中使用 H1Loss 训练、LpLoss 评估。文档明确说明,H1Loss 同时惩罚函数值和梯度误差,因此适合 PDE 问题。(Neural Operator)
立即加上以下检查:
python
batch = next(iter(train_loader))
print(batch.keys())
print("x:", batch["x"].shape)
print("y:", batch["y"].shape)
with torch.no_grad():
x = batch["x"].to(device)
y_pred = model(x)
print("prediction:", y_pred.shape)你的第一目标不是调高精度,而是能够回答:
输入物理量是什么?
输出物理量是什么?
tensor 每一维表示什么?
模型如何从输入场产生输出场?
在不同网格分辨率上,模型行为有何变化?
4. 第二天:扔掉 Trainer,手写普通 PyTorch 训练循环
只会运行官方封装不算真正上手。下一步是自己写训练循环。
python
for epoch in range(20):
model.train()
train_loss = 0.0
for batch in train_loader:
batch = data_processor.preprocess(batch)
x = batch["x"].to(device)
y = batch["y"].to(device)
optimizer.zero_grad()
y_pred = model(x)
loss = h1_loss(y_pred, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
scheduler.step()
print(
f"epoch={epoch:02d}, "
f"loss={train_loss / len(train_loader):.6f}"
)然后自己写验证函数:
python
@torch.no_grad()
def evaluate(model, loader, data_processor, loss_fn, device):
model.eval()
total_loss = 0.0
for batch in loader:
batch = data_processor.preprocess(batch)
x = batch["x"].to(device)
y = batch["y"].to(device)
y_pred = model(x)
total_loss += loss_fn(y_pred, y).item()
return total_loss / len(loader)此时你已经把 Neural Operator 降维成一个熟悉的问题:
text
Dataset
→ DataLoader
→ model(x)
→ loss(pred, y)
→ backward()
→ evaluate()后续换 PDE、换架构、加入物理约束,本质上都在修改这个框架。
5. 第三步:自己实现一个极简 FNO
只调用库无法建立结构直觉。你至少应该手写一次二维 Spectral Convolution。
FNO 的核心不是普通卷积,而是:
对空间场做 FFT;
保留部分低频 Fourier modes;
在频域进行可学习线性变换;
做逆 FFT 回到空间域;
与局部线性分支结合。
一个教学用最小实现如下:
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class SpectralConv2d(nn.Module):
def __init__(self, in_channels, out_channels, modes_x, modes_y):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.modes_x = modes_x
self.modes_y = modes_y
scale = 1.0 / (in_channels * out_channels)
self.weight = nn.Parameter(
scale
* torch.randn(
in_channels,
out_channels,
modes_x,
modes_y,
dtype=torch.cfloat,
)
)
def forward(self, x):
batch_size, _, height, width = x.shape
x_ft = torch.fft.rfft2(x)
out_ft = torch.zeros(
batch_size,
self.out_channels,
height,
width // 2 + 1,
dtype=torch.cfloat,
device=x.device,
)
out_ft[:, :, : self.modes_x, : self.modes_y] = torch.einsum(
"bixy,ioxy->boxy",
x_ft[:, :, : self.modes_x, : self.modes_y],
self.weight,
)
return torch.fft.irfft2(
out_ft,
s=(height, width),
)
class SimpleFNO2d(nn.Module):
def __init__(self, modes_x=12, modes_y=12, width=32):
super().__init__()
self.lift = nn.Conv2d(1, width, kernel_size=1)
self.spectral_layers = nn.ModuleList(
[
SpectralConv2d(width, width, modes_x, modes_y)
for _ in range(4)
]
)
self.pointwise_layers = nn.ModuleList(
[
nn.Conv2d(width, width, kernel_size=1)
for _ in range(4)
]
)
self.project = nn.Sequential(
nn.Conv2d(width, 64, kernel_size=1),
nn.GELU(),
nn.Conv2d(64, 1, kernel_size=1),
)
def forward(self, x):
x = self.lift(x)
for spectral, pointwise in zip(
self.spectral_layers,
self.pointwise_layers,
):
x = F.gelu(spectral(x) + pointwise(x))
return self.project(x)这个版本刻意省略了不少工程细节,例如:
正负频率的完整处理;
positional embedding;
padding;
normalization;
residual 设计;
factorized weights;
多维和复杂边界情况。
但它足以回答最关键的问题:
FNO 如何用 FFT 将全局空间耦合转换为频域中的可学习变换?
然后做一个 sanity check:
python
model = SimpleFNO2d()
x = torch.randn(8, 1, 64, 64)
y = model(x)
print(x.shape)
print(y.shape)预期:
text
torch.Size([8, 1, 64, 64])
torch.Size([8, 1, 64, 64])6. 应按什么顺序做项目?
不要直接从 Navier–Stokes 起步。复杂时序 rollout、数值稳定性和数据生成会同时干扰你。
推荐顺序
| 顺序 | 任务 | 目的 |
|---|---|---|
| 1 | Darcy Flow 2D | 学会输入场到输出场的静态映射 |
| 2 | Burgers Equation 1D | 学习时序场、autoregressive rollout |
| 3 | Diffusion / Advection | 理解平滑、传播和频域行为 |
| 4 | Darcy Flow:自己生成数据 | 接触数值求解器和数据管线 |
| 5 | Navier–Stokes 2D | 做相对完整的复现项目 |
| 6 | PDEBench 子任务 | 进入规范化 benchmark |
官方 NeuralOperator 文档已经提供 Darcy Flow 数据集、FNO 训练示例,以及 diffusion-advection、二维 Burgers 数据生成示例。(Neural Operator)
PDEBench 则提供更广泛的时变 PDE benchmark、数据生成代码、ready-to-use 数据和 FNO、U-Net、PINN 等 baseline。它适合第二阶段,不适合作为第一次接触 Neural Operator 的起点。(GitHub)
7. DeepONet 什么时候学?
在 FNO 跑通后,用半天时间实现 DeepONet。
DeepONet 的结构非常直观:
text
输入函数在若干传感器位置上的采样值
↓
branch net
查询位置坐标 (x, y, t)
↓
trunk net
两者做内积
↓
输出 u(x, y, t)DeepONet 原始论文将其描述为两个子网络:
branch net:编码输入函数在固定 sensor 上的取值;trunk net:编码待查询位置;输出由二者组合得到。(arXiv)
极简实现:
python
import torch
import torch.nn as nn
class MLP(nn.Module):
def __init__(self, dims):
super().__init__()
layers = []
for in_dim, out_dim in zip(dims[:-2], dims[1:-1]):
layers.append(nn.Linear(in_dim, out_dim))
layers.append(nn.GELU())
layers.append(nn.Linear(dims[-2], dims[-1]))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
class DeepONet(nn.Module):
def __init__(self, n_sensors, coord_dim, hidden_dim=128):
super().__init__()
self.branch = MLP(
[n_sensors, hidden_dim, hidden_dim, hidden_dim]
)
self.trunk = MLP(
[coord_dim, hidden_dim, hidden_dim, hidden_dim]
)
def forward(self, sensor_values, coords):
branch_features = self.branch(sensor_values)
trunk_features = self.trunk(coords)
return torch.sum(
branch_features * trunk_features,
dim=-1,
keepdim=True,
)理解 FNO 和 DeepONet 后,你会自然获得一个重要区分:
| 架构 | 更自然的输入形式 | 典型优势 |
|---|---|---|
| FNO | 规则网格上的场 | 图像式 PDE 数据、FFT 高效 |
| DeepONet | 函数采样值 + 查询坐标 | 查询位置灵活、结构直观 |
| GNO / GINO | 点云、mesh、几何结构 | 不规则几何和工程场景 |
8. 一个合理的 7 天冲刺计划
Day 1:跑通官方示例
完成:
text
安装 neuraloperator
运行 Darcy Flow FNO
输出 x/y/prediction shape
确认 GPU 可用
记录 16×16 和 32×32 的测试误差Day 2:可视化与实验记录
画出:
text
input permeability
ground truth pressure
predicted pressure
absolute error修改:
text
n_modes = 4, 8, 12
hidden_channels = 16, 24, 48
n_train = 100, 500, 1000观察模型容量、频域截断和数据量对误差的影响。
Day 3:手写训练循环
去掉 Trainer,自己实现:
text
train_one_epoch()
evaluate()
save_checkpoint()
load_checkpoint()
plot_prediction()Day 4:手写 SpectralConv2d
实现极简 FNO,跑 shape test,再替换官方 FNO 做对比。
Day 5:自己生成简单数据
选一个低难度 PDE:
text
1D Burgers
1D diffusion
2D diffusion-advection将数据保存为:
python
x.shape == [num_samples, channels, grid_size]
y.shape == [num_samples, channels, grid_size]Day 6:实现 DeepONet
用同一个 PDE 数据集比较:
text
MLP baseline
DeepONet
FNODay 7:整理为公开仓库
仓库结构:
text
neural-operator-starter/
├── README.md
├── requirements.txt
├── configs/
│ ├── darcy_fno.yaml
│ └── burgers_fno.yaml
├── src/
│ ├── models/
│ │ ├── simple_fno.py
│ │ └── deeponet.py
│ ├── train.py
│ ├── evaluate.py
│ └── visualize.py
├── notebooks/
│ └── 01_darcy_fno_walkthrough.ipynb
└── results/
├── figures/
└── metrics.csvREADME 中至少写清楚:
text
Problem
Dataset
Tensor shapes
Model
Loss
Training setup
Results
Failure cases
How to reproduce9. 暂时不要做什么
初学阶段不要立即投入:
复杂 CFD 工业数据;
非规则 mesh;
多 GPU 分布式训练;
PINO、physics-informed loss;
TFNO、SFNO、GINO、U-NO 等大量变体;
一开始就复现最新 SOTA;
同时学习完整 PDE 理论、数值分析、泛函分析和湍流理论。
这些方向后续有价值,但会显著增加变量数量。当前目标是建立闭环:
text
理解数据
→ 跑通模型
→ 手写训练
→ 手写核心层
→ 比较 baseline
→ 形成公开仓库10. 最值得直接使用的资料
优先级从高到低:
官方 NeuralOperator 仓库:安装、API、源码。(GitHub)
官方 FNO Darcy Flow 教程:最小可运行闭环,包含训练、可视化和 zero-shot super-resolution。(Neural Operator)
官方 examples gallery:后续切换数据生成、Burgers、SFNO、U-NO、OTNO。(Neural Operator)
Zongyi Li 的 FNO 介绍文章:用于快速建立直觉。(Zongyi Li)
PDEBench:完成入门后进入 benchmark。(GitHub)
DeepONet 原始仓库与论文:用于补充另一类 operator learning 范式。(GitHub)
最实际的起点是:今天只完成官方 Darcy Flow 示例、打印 tensor shape、画出预测图,并将 Trainer 替换为自己的 PyTorch 循环。这样一天内就能从“知道 Neural Operator”进入“能改 Neural Operator 代码”。