跳转至

neuraloperator

# darcy-flow 数据集下载
# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_train_16.npy -P ./datasets/darcyflow/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_32.npy -P ./datasets/darcyflow/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_16.npy -P ./datasets/darcyflow/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_train_16.npy -o ./datasets/darcyflow/darcy_train_16.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_32.npy -o ./datasets/darcyflow/darcy_test_32.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_16.npy -o ./datasets/darcyflow/darcy_test_16.npy
# tfno 模型训练
python train_tfno.py
# uno 模型训练
python train_uno.py

# SWE 数据集下载
# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/train_SWE_32x64.npy -P ./datasets/SWE/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_64x128.npy -P ./datasets/SWE/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_32x64.npy -P ./datasets/SWE/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/train_SWE_32x64.npy -o ./datasets/SWE/train_SWE_32x64.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_64x128.npy -o ./datasets/SWE/test_SWE_64x128.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_32x64.npy -o ./datasets/SWE/test_SWE_32x64.npy

# sfno 模型训练
python train_sfno.py

# darcy-flow 数据集下载
# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_train_16.npy -P ./datasets/darcyflow/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_32.npy -P ./datasets/darcyflow/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_16.npy -P ./datasets/darcyflow/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_train_16.npy -o ./datasets/darcyflow/darcy_train_16.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_32.npy -o ./datasets/darcyflow/darcy_test_32.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/darcy_flow/darcy_test_16.npy -o ./datasets/darcyflow/darcy_test_16.npy
# tfno 模型评估
python train_tfno.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_tfno.pdparams
# uno 模型评估
python train_uno.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_uno.pdparams

# SWE 数据集下载
# linux
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/train_SWE_32x64.npy -P ./datasets/SWE/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_64x128.npy -P ./datasets/SWE/
wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_32x64.npy -P ./datasets/SWE/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/train_SWE_32x64.npy -o ./datasets/SWE/train_SWE_32x64.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_64x128.npy -o ./datasets/SWE/test_SWE_64x128.npy
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/neuraloperator/SWE_data/test_SWE_32x64.npy -o ./datasets/SWE/test_SWE_32x64.npy
# sfno 模型评估
python train_sfno.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_sfno.pdparams

# tfno 模型导出
python train_tfno.py mode=export INFER.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_tfno.pdparams
# uno 模型导出
python train_uno.py mode=export INFER.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_uno.pdparams
# sfno 模型导出
python train_sfno.py mode=export INFER.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/neuraloperator/neuraloperator_sfno.pdparams

# tfno 模型推理
python train_tfno.py mode=infer
# uno 模型推理
python train_uno.py mode=infer
# sfno 模型推理
python train_sfno.py mode=infer
模型 16_h1 16_l2 32_h1 32_l2
tfno 模型 0.13113 0.08514 0.30353 0.12408
模型 16_h1 16_l2 32_h1 32_l2
uno 模型 0.18360 0.11040 0.74840 0.60193
模型 32x64_l2 64x128_l2
sfno 模型 1.01075 2.33481

1. 背景简介

许多科学和工程问题，如分子动力学、微力学和湍流流动，都需要反复求解复杂的偏微分方程（PDE）系统，以便获取某些参数的不同值。为了准确捕捉所模拟的现象，这些系统通常需要进行精细的离散化。然而，这也导致了传统数值求解器运行缓慢，有时甚至效率低下。在这种情况下，机器学习方法有望通过提供快速求解器来革新科学领域，这些求解器能够近似或增强传统方法。但值得注意的是，经典神经网络是在有限维空间之间进行映射，因此它们只能学习与特定离散化相关的解决方案，这在实际应用中是一个限制。为了克服这一限制，最近的一项新研究提出了使用神经网络来学习无网格、无限维的算子。这种神经算子通过生成一组用于不同离散化且与网格无关的参数，解决了有限维算子方法中的网格依赖性问题。该研究通过直接在傅里叶空间中参数化积分核，制定了一个新的神经算子，从而创建了一个富有表现力和高效的架构。论文中对 Burgers 方程、Darcy 流和 Navier-Stokes 方程进行了实验验证。值得一提的是，傅里叶神经算子作为首个基于机器学习的方法，成功地以零样本超分辨率模拟了湍流，其速度比传统PDE求解器快达三个数量级。

2. 模型原理

本章节仅对 NeuralOperator 的模型原理进行简单地介绍，详细的理论推导请阅读Fourier Neural Operator for Parametric Partial Differential Equations。 NeuralOperator 引入了傅里叶神经算子 (Fourier neural operator)，这是一种新颖的深度学习架构，能够学习函数之间无限维空间的映射；积分算子被限制为卷积，并通过傅里叶域中的线性变换实例化。傅里叶神经算子是第一个学习湍流状态下 Navier-Stokes 方程族的分辨率不变解算子的工作，其中以前基于图形的神经算子不收敛。该方法共享相同的学习网络参数，而不考虑输入和输出空间上使用的离散化。

模型的总体结构如图所示：

NeuralOperator-arch

NeuralOperator 网络模型

NeuralOperator 论文中使用 TFNO 和 UNO 模型训练 Darcy-Flow 数据集，并进行验证和推理；使用 SFNO 模型训练 Spherical Shallow Water(SWE) 数据集，并进行验证和推理。接下来分别进行介绍。

2.1 模型训练、推理过程

模型预训练阶段是基于随机初始化的网络权重对模型进行训练，如下图所示，其中 \(X_{[w,h]}\) 表示大小为 \(w*h\) 的二维偏微分数据，\(Y_{[w,h]}\) 表示预测的大小为 \(w*h\) 的二维偏微分方程数值解，\(Y_{true[w,h]}\) 表示真实二维偏微分方程数值解。最后网络模型预测的输出和真值计算 LpLoss 或者 H1 损失函数。

FNO-pretraining

FNO 模型预训练

在推理阶段，给定大小为 \(w*h\) 的二维偏微分数据，预测得到大小为 \(w*h\) 的二维偏微分方程数值解。

FNO-infer

FNO 模型推理

3. TFNO 模型训练 darcy-flow 实现

接下来开始讲解如何基于 PaddleScience 代码，实现 TFNO 模型对 darcy-flow 数据的训练与推理。关于该案例中的其余细节请参考 API文档。

3.1 数据集介绍

使用 二维达西流 (darcy-flow) 数据集，这个问题的偏微分方程为：

\(-\nabla\cdot (k(x)\nabla u(x))=f(x),x\in D\)

其中，x 是位置，u(x) 是流体的压力，k(x) 是渗透率场，f(x) 是压力的函数。达西流问题可以被用来描述多孔介质的流动、弹性材料和热传导。在这里，我们定义了一个二维的平面区域 \(D=[0,1]×[0,1]\)，我们希望得到一个模型，可以在给定 k 渗透率场的情况下，估算出 u 流体压力。

训练数据和测试数据：

数据集包括 1000 条 16x16 分辨率大小的训练数据；50 条 32x32 和 50 条 32x32 分辨率大小的测试数据。数据格式采用 NPY 格式保存。

3.2 模型预训练

3.2.1 约束构建

本案例基于数据驱动的方法求解问题，因此需要使用 PaddleScience 内置的 SupervisedConstraint 构建监督约束。在定义约束之前，需要首先指定监督约束中用于数据加载的各个参数。

数据加载的代码如下:

examples/neuraloperator/train_tfno.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# set train dataloader config
train_dataloader_cfg = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "train",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
    "num_workers": 0,
}

其中，"dataset" 字段定义了使用的 Dataset 类名为 DarcyFlowDataset，"sampler" 字段定义了使用的 Sampler 类名为 BatchSampler，设置的 batch_size 为 16，num_works 为 0。

定义监督约束的代码如下：

examples/neuraloperator/train_tfno.py
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# set loss
l2loss = metric.LpLoss_train(d=2, p=2)
h1loss = metric.H1Loss_train(d=2)
if cfg.TRAIN.training_loss == "l2":
    train_loss = l2loss
if cfg.TRAIN.training_loss == "h1":
    train_loss = h1loss

# set constraint
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    name="Sup",
)
constraint = {sup_constraint.name: sup_constraint}

SupervisedConstraint 的第一个参数是数据的加载方式，这里使用上文中定义的 train_dataloader_cfg；

第二个参数是损失函数的定义，这里使用自定义的损失函数 h1；

第三个参数是约束条件的名字，方便后续对其索引。此处命名为 Sup。

3.2.2 模型构建

在该案例中，darcy-flow 基于 TFNO 网络模型实现，用 PaddleScience 代码表示如下：

examples/neuraloperator/train_tfno.py
131
132
133
model = ppsci.arch.TFNO2dNet(
    **cfg.MODEL,
)

网络模型的参数通过配置文件进行设置如下：

examples/neuraloperator/conf/tfno_darcyflow_pretrain.yaml
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
# model settings
MODEL:
  input_keys: ["x"]
  output_keys: ["y"]
  n_modes_height: 16
  n_modes_width: 16
  in_channels: 3
  out_channels: 1
  hidden_channels: 32
  projection_channels: 64
  n_layers: 4

  use_mlp: False
  mlp:
    expansion: 0.5
    dropout: 0.0
  norm: "group_norm"
  fno_skip: "linear"
  mlp_skip: "soft-gating"
  separable: false
  preactivation: false
  factorization: "dense"
  rank: 1.0
  joint_factorization: false
  fixed_rank_modes: null
  implementation: "factorized"
  domain_padding: null #0.078125
  domain_padding_mode: "one-sided" #symmetric
  fft_norm: "forward"

其中，input_keys 和 output_keys 分别代表网络模型输入、输出变量的名称。

3.2.3 学习率与优化器构建

本案例中使用的学习率方法为 StepDecay，学习率大小设置为 5e-3。优化器使用 Adam,用 PaddleScience 代码表示如下：

examples/neuraloperator/train_tfno.py
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# init optimizer and lr scheduler
if cfg.TRAIN.lr_scheduler.type == "ReduceOnPlateau":
    lr_scheduler = paddle.optimizer.lr.ReduceOnPlateau(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        factor=cfg.TRAIN.lr_scheduler.gamma,
        patience=cfg.TRAIN.lr_scheduler.scheduler_patience,
        mode="min",
    )
elif cfg.TRAIN.lr_scheduler.type == "CosineAnnealingDecay":
    lr_scheduler = paddle.optimizer.lr.CosineAnnealingDecay(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        T_max=cfg.TRAIN.lr_scheduler.scheduler_T_max,
    )()
elif cfg.TRAIN.lr_scheduler.type == "StepDecay":
    lr_scheduler = ppsci.optimizer.lr_scheduler.Step(
        epochs=cfg.TRAIN.lr_scheduler.epochs,
        iters_per_epoch=ITERS_PER_EPOCH,
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        step_size=cfg.TRAIN.lr_scheduler.step_size,
        gamma=cfg.TRAIN.lr_scheduler.gamma,
    )()
else:
    raise ValueError(f"Got scheduler={cfg.TRAIN.lr_scheduler.type}")
optimizer = ppsci.optimizer.Adam(lr_scheduler, weight_decay=cfg.TRAIN.wd)(model)

3.2.4 评估器构建

本案例训练过程中会按照一定的训练轮数间隔，使用验证集评估当前模型的训练情况，需要使用 SupervisedValidator 构建评估器。代码如下：

examples/neuraloperator/train_tfno.py
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
# set eval dataloader config
eval_dataloader_cfg_16 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_16x16",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_32x32",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

h1_eval_metric = metric.H1Loss(d=2)
l2_eval_metric = metric.LpLoss(d=2, p=2)
sup_validator_16 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_16,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_16x16",
)

sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_32x32",
)

validator = {
    sup_validator_16.name: sup_validator_16,
    sup_validator_32.name: sup_validator_32,
}

SupervisedValidator 评估器与 SupervisedConstraint 比较相似，不同的是评估器需要设置评价指标 metric，在这里使用了自定义的评价指标分别是 hlLoss 和 LpLoss。

3.2.5 模型训练与评估

完成上述设置之后，只需要将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动训练、评估。

examples/neuraloperator/train_tfno.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    lr_scheduler,
    cfg.TRAIN.epochs,
    ITERS_PER_EPOCH,
    eval_during_train=cfg.TRAIN.eval_during_train,
    seed=cfg.seed,
    validator=validator,
    compute_metric_by_batch=cfg.EVAL.compute_metric_by_batch,
    eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
    pretrained_model_path=cfg.TRAIN.pretrained_model_path,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

3.3 模型评估可视化

3.3.1 测试集上评估模型

构建模型的代码为：

examples/neuraloperator/train_tfno.py
265
266
267
model = ppsci.arch.TFNO2dNet(
    **cfg.MODEL,
)

构建评估器的代码为：

examples/neuraloperator/train_tfno.py
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# set eval dataloader config
eval_dataloader_cfg_16 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_16x16",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_32x32",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

# set loss
l2loss = metric.LpLoss_train(d=2, p=2)
h1loss = metric.H1Loss_train(d=2)
if cfg.TRAIN.training_loss == "l2":
    train_loss = l2loss
if cfg.TRAIN.training_loss == "h1":
    train_loss = h1loss

h1_eval_metric = metric.H1Loss(d=2)
l2_eval_metric = metric.LpLoss(d=2, p=2)
sup_validator_16 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_16,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_16x16",
)

sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_32x32",
)
validator = {
    sup_validator_16.name: sup_validator_16,
    sup_validator_32.name: sup_validator_32,
}

3.3.2 模型导出

构建模型的代码为：

examples/neuraloperator/train_tfno.py
285
286
287
288
# set model
model = ppsci.arch.TFNO2dNet(
    **cfg.MODEL,
)

实例化 ppsci.solver.Solver：

examples/neuraloperator/train_tfno.py
290
291
292
293
294
# initialize solver
solver = ppsci.solver.Solver(
    model,
    pretrained_model_path=cfg.INFER.pretrained_model_path,
)

构建模型输入格式并导出静态模型：

examples/neuraloperator/train_tfno.py
295
296
297
298
299
300
301
302
303
304
# export model
from paddle.static import InputSpec

input_spec = [
    {
        key: InputSpec([None, 3, 16, 16], "float32", name=key)
        for key in model.input_keys
    },
]
solver.export(input_spec, cfg.INFER.export_path)

InputSpec 函数中第一个设置模型输入尺寸，第二个参数设置输入数据类型，第三个设置输入数据的 Key.

3.3.3 模型推理

创建预测器:

examples/neuraloperator/train_tfno.py
309
310
311
import predictor

predictor = predictor.FNOPredictor(cfg)

准备预测数据：

examples/neuraloperator/train_tfno.py
313
314
315
316
data = np.load(cfg.INFER.data_path, allow_pickle=True).item()

input_data = data["x"][0].reshape(-1, 1, *data["x"].shape[1:]).astype("float32")
label = data["y"][0].astype("float32")

进行模型预测与预测值显示:

examples/neuraloperator/train_tfno.py
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
pred_data = predictor.predict(input_data, cfg.INFER.batch_size)

fig = plt.figure(figsize=(7, 7))

ax = fig.add_subplot(1, 3, 1)
ax.imshow(input_data.squeeze(), cmap="gray")
ax.set_title("k(x)")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 2)
ax.imshow(label)
ax.set_title("Ground-truth y")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 3)
ax.imshow(pred_data.squeeze())
ax.set_title("Model prediction")
plt.xticks([], [])
plt.yticks([], [])
plt.savefig(cfg.output_dir)
logger.message("save success")
plt.close(fig)

4. UNO 模型训练 darcy-flow 实现

4.1 数据集介绍

数据集同 3.1 节。

4.2 模型预训练

4.2.1 约束构建

本案例基于数据驱动的方法求解问题，因此需要使用 PaddleScience 内置的 SupervisedConstraint 构建监督约束。在定义约束之前，需要首先指定监督约束中用于数据加载的各个参数。

数据加载的代码如下:

examples/neuraloperator/train_uno.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# set train dataloader config
train_dataloader_cfg = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "train",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": True,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
    "num_workers": 0,
}

其中，"dataset" 字段定义了使用的 Dataset 类名为 DarcyFlowDataset，"sampler" 字段定义了使用的 Sampler 类名为 BatchSampler，设置的 batch_size 为 16，num_works 为 0。

定义监督约束的代码如下：

examples/neuraloperator/train_uno.py
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# set loss
l2loss = metric.LpLoss_train(d=2, p=2)
h1loss = metric.H1Loss_train(d=2)
if cfg.TRAIN.training_loss == "l2":
    train_loss = l2loss
if cfg.TRAIN.training_loss == "h1":
    train_loss = h1loss

# set constraint
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    name="Sup",
)
constraint = {sup_constraint.name: sup_constraint}

SupervisedConstraint 的第一个参数是数据的加载方式，这里使用上文中定义的 train_dataloader_cfg；

第二个参数是损失函数的定义，这里使用自定义的损失函数 h1；

第三个参数是约束条件的名字，方便后续对其索引。此处命名为 Sup。

4.2.2 模型构建

在该案例中，darcy-flow 基于 UNO 网络模型实现，用 PaddleScience 代码表示如下：

examples/neuraloperator/train_uno.py
131
132
133
model = ppsci.arch.UNONet(
    **cfg.MODEL,
)

网络模型的参数通过配置文件进行设置如下：

examples/neuraloperator/conf/uno_darcyflow_pretrain.yaml
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# model settings
MODEL:
  input_keys: ["x"]
  output_keys: ["y"]
  in_channels: 3
  out_channels: 1
  hidden_channels: 64
  projection_channels: 64
  n_layers: 5
  uno_out_channels: [32, 64, 64, 64, 32]
  uno_n_modes: [[16, 16], [8, 8], [8, 8], [8, 8], [16, 16]]
  uno_scalings: [[1.0, 1.0], [0.5, 0.5], [1, 1], [2, 2], [1, 1]]
  horizontal_skips_map: null
  incremental_n_modes: null

  use_mlp: false
  mlp:
    expansion: 0.5
    dropout: 0.0
  norm: "group_norm"
  fno_skip: "linear"
  horizontal_skip: "linear"
  mlp_skip: "soft-gating"
  separable: false
  preactivation: false
  factorization: null
  rank: 1.0
  joint_factorization: false
  fixed_rank_modes: null
  implementation: "factorized"
  domain_padding: 0.2 #0.078125
  domain_padding_mode: "one-sided" #symmetric
  fft_norm: "forward"

其中，input_keys 和 output_keys 分别代表网络模型输入、输出变量的名称。

4.2.3 学习率与优化器构建

本案例中使用的学习率方法为 StepDecay，学习率大小设置为 5e-3。优化器使用 Adam,用 PaddleScience 代码表示如下：

examples/neuraloperator/train_uno.py
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# init optimizer and lr scheduler
if cfg.TRAIN.lr_scheduler.type == "ReduceOnPlateau":
    lr_scheduler = paddle.optimizer.lr.ReduceOnPlateau(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        factor=cfg.TRAIN.lr_scheduler.gamma,
        patience=cfg.TRAIN.lr_scheduler.scheduler_patience,
        mode="min",
    )
elif cfg.TRAIN.lr_scheduler.type == "CosineAnnealingDecay":
    lr_scheduler = paddle.optimizer.lr.CosineAnnealingDecay(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        T_max=cfg.TRAIN.lr_scheduler.scheduler_T_max,
    )()
elif cfg.TRAIN.lr_scheduler.type == "StepDecay":
    lr_scheduler = ppsci.optimizer.lr_scheduler.Step(
        epochs=cfg.TRAIN.lr_scheduler.epochs,
        iters_per_epoch=ITERS_PER_EPOCH,
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        step_size=cfg.TRAIN.lr_scheduler.step_size,
        gamma=cfg.TRAIN.lr_scheduler.gamma,
    )()
else:
    raise ValueError(f"Got scheduler={cfg.TRAIN.lr_scheduler.type}")
optimizer = ppsci.optimizer.Adam(lr_scheduler, weight_decay=cfg.TRAIN.wd)(model)

4.2.4 评估器构建

本案例训练过程中会按照一定的训练轮数间隔，使用验证集评估当前模型的训练情况，需要使用 SupervisedValidator 构建评估器。代码如下：

examples/neuraloperator/train_uno.py
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
# set eval dataloader config
eval_dataloader_cfg_16 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_16x16",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_32x32",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

h1_eval_metric = metric.H1Loss(d=2)
l2_eval_metric = metric.LpLoss(d=2, p=2)
sup_validator_16 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_16,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_16x16",
)

sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_32x32",
)

validator = {
    sup_validator_16.name: sup_validator_16,
    sup_validator_32.name: sup_validator_32,
}

SupervisedValidator 评估器与 SupervisedConstraint 比较相似，不同的是评估器需要设置评价指标 metric，在这里使用了自定义的评价指标分别是 hlLoss 和 LpLoss。

4.2.5 模型训练与评估

完成上述设置之后，只需要将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动训练、评估。

examples/neuraloperator/train_uno.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    lr_scheduler,
    cfg.TRAIN.epochs,
    ITERS_PER_EPOCH,
    eval_during_train=cfg.TRAIN.eval_during_train,
    seed=cfg.seed,
    validator=validator,
    compute_metric_by_batch=cfg.EVAL.compute_metric_by_batch,
    eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
    pretrained_model_path=cfg.TRAIN.pretrained_model_path,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

4.3 模型评估可视化

4.3.1 测试集上评估模型

构建模型的代码为：

examples/neuraloperator/train_uno.py
265
266
267
model = ppsci.arch.UNONet(
    **cfg.MODEL,
)

构建评估器的代码为：

examples/neuraloperator/train_uno.py
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# set eval dataloader config
eval_dataloader_cfg_16 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_16x16",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "DarcyFlowDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "grid_boundaries": cfg.DATASET.grid_boundaries,
        "encode_input": cfg.DATASET.encode_input,
        "encode_output": cfg.DATASET.encode_output,
        "encoding": cfg.DATASET.encoding,
        "channel_dim": cfg.DATASET.channel_dim,
        "data_split": "test_32x32",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

# set loss
l2loss = metric.LpLoss_train(d=2, p=2)
h1loss = metric.H1Loss_train(d=2)
if cfg.TRAIN.training_loss == "l2":
    train_loss = l2loss
if cfg.TRAIN.training_loss == "h1":
    train_loss = h1loss

h1_eval_metric = metric.H1Loss(d=2)
l2_eval_metric = metric.LpLoss(d=2, p=2)
sup_validator_16 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_16,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_16x16",
)

sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={
        "h1": ppsci.metric.FunctionalMetric(h1_eval_metric),
        "l2": ppsci.metric.FunctionalMetric(l2_eval_metric),
    },
    name="Sup_Validator_32x32",
)
validator = {
    sup_validator_16.name: sup_validator_16,
    sup_validator_32.name: sup_validator_32,
}

4.3.2 模型导出

构建模型的代码为：

examples/neuraloperator/train_uno.py
285
286
287
288
# set model
model = ppsci.arch.UNONet(
    **cfg.MODEL,
)

实例化 ppsci.solver.Solver：

examples/neuraloperator/train_uno.py
290
291
292
293
294
# initialize solver
solver = ppsci.solver.Solver(
    model,
    pretrained_model_path=cfg.INFER.pretrained_model_path,
)

构建模型输入格式并导出静态模型：

examples/neuraloperator/train_uno.py
295
296
297
298
299
300
301
302
303
304
# export model
from paddle.static import InputSpec

input_spec = [
    {
        key: InputSpec([None, 3, 16, 16], "float32", name=key)
        for key in model.input_keys
    },
]
solver.export(input_spec, cfg.INFER.export_path)

InputSpec 函数中第一个设置模型输入尺寸，第二个参数设置输入数据类型，第三个设置输入数据的 Key.

4.3.3 模型推理

创建预测器:

examples/neuraloperator/train_uno.py
309
310
311
import predictor

predictor = predictor.FNOPredictor(cfg)

准备预测数据：

examples/neuraloperator/train_uno.py
313
314
315
316
data = np.load(cfg.INFER.data_path, allow_pickle=True).item()

input_data = data["x"][0].reshape(-1, 1, *data["x"].shape[1:]).astype("float32")
label = data["y"][0].astype("float32")

进行模型预测与预测值显示:

examples/neuraloperator/train_uno.py
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
pred_data = predictor.predict(input_data, cfg.INFER.batch_size)

fig = plt.figure(figsize=(7, 7))

ax = fig.add_subplot(1, 3, 1)
ax.imshow(input_data.squeeze(), cmap="gray")
ax.set_title("k(x)")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 2)
ax.imshow(label)
ax.set_title("Ground-truth y")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 3)
ax.imshow(pred_data.squeeze())
ax.set_title("Model prediction")
plt.xticks([], [])
plt.yticks([], [])
plt.savefig(cfg.output_dir)
logger.message("save success")
plt.close(fig)

5. SFNO 模型训练 spherical Shallow Water equations(SWE) 实现

5.1 数据集介绍

球面浅水方程（Spherical Shallow Water Equations，简称 SWE）是一组描述在旋转地球表面上的浅水流动的偏微分方程。浅水方程通常用于模拟海洋、湖泊和河流中的流体运动，当流体的垂直尺度远小于其水平尺度时，可以忽略流体的垂直结构，只考虑其水平运动。

球面浅水方程在数学上可以由以下方程组表示：

\(\frac{\partial u}{\partial t} +u\cdot \nabla u=-g\nabla h-fu+F\)

\(\frac{\partial h}{\partial t}+\nabla \cdot (hu)=0\)

其中：

𝑢 是水平速度场，通常包含经度和纬度方向的速度分量。

ℎ 是流体高度（或水面高度）相对于参考水平面的位移。

𝑔 是重力加速度。

𝑓 是科里奥利参数，它与地球自转和纬度有关，f=2Ωsinϕ，其中 Ω 是地球自转的角度，𝜙 是纬度。

𝐹 是摩擦力和其他外部力（如风力）的向量。

∇ 是水平梯度算子。

球面浅水方程考虑了地球的球形几何，因此使用的是球面坐标系。在实际应用中，这些方程通常需要进行离散化和数值求解，以便于在计算机上进行模拟。

训练数据和测试数据：

数据集包括 200 条 32x64 分辨率大小的训练数据；50 条 32x64 和 50 条 64x128 分辨率大小的测试数据。数据格式采用 NPY 格式保存。

5.2 模型预训练

5.2.1 约束构建

本案例基于数据驱动的方法求解问题，因此需要使用 PaddleScience 内置的 SupervisedConstraint 构建监督约束。在定义约束之前，需要首先指定监督约束中用于数据加载的各个参数。

数据加载的代码如下:

examples/neuraloperator/train_sfno.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# set train dataloader config
train_dataloader_cfg = {
    "dataset": {
        "name": "SphericalSWEDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "data_split": "train",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
    "num_workers": 0,
}

其中，"dataset" 字段定义了使用的 Dataset 类名为 DarcyFlowDataset，"sampler" 字段定义了使用的 Sampler 类名为 BatchSampler，设置的 batch_size 为 4，num_works 为 0。

定义监督约束的代码如下：

examples/neuraloperator/train_sfno.py
32
33
34
35
36
37
38
39
40
41
# set loss
train_loss = metric.LpLoss_train(d=2, p=2, reduce_dims=[0, 1])

# set constraint
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    name="Sup",
)
constraint = {sup_constraint.name: sup_constraint}

SupervisedConstraint 的第一个参数是数据的加载方式，这里使用上文中定义的 train_dataloader_cfg；

第二个参数是损失函数的定义，这里使用自定义的损失函数 Lp；

第三个参数是约束条件的名字，方便后续对其索引。此处命名为 Sup。

5.2.2 模型构建

在该案例中，SWE 基于 SFNO 网络模型实现，用 PaddleScience 代码表示如下：

examples/neuraloperator/train_sfno.py
104
105
106
model = ppsci.arch.SFNONet(
    **cfg.MODEL,
)

网络模型的参数通过配置文件进行设置如下：

examples/neuraloperator/conf/sfno_swe_pretrain.yaml
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
# model settings
MODEL:
  input_keys: ["x"]
  output_keys: ["y"]
  in_channels: 3
  out_channels: 3
  n_modes: [32, 32]
  hidden_channels: 32
  projection_channels: 64
  n_layers: 4

  use_mlp: false
  mlp:
    expansion: 0.5
    dropout: 0.0
  norm: 'group_norm'
  fno_skip: "linear"
  mlp_skip: "soft-gating"
  separable: false
  preactivation: false
  factorization: null
  rank: 1.0
  joint_factorization: false
  fixed_rank_modes: null
  implementation: "factorized"
  domain_padding: null #0.078125
  domain_padding_mode: "one-sided" #symmetric
  fft_norm: 'forward'

其中，input_keys 和 output_keys 分别代表网络模型输入、输出变量的名称。

5.2.3 学习率与优化器构建

本案例中使用的学习率方法为 StepDecay，学习率大小设置为 5e-3。优化器使用 Adam,用 PaddleScience 代码表示如下：

examples/neuraloperator/train_sfno.py
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# init optimizer and lr scheduler
if cfg.TRAIN.lr_scheduler.type == "ReduceOnPlateau":
    lr_scheduler = paddle.optimizer.lr.ReduceOnPlateau(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        factor=cfg.TRAIN.lr_scheduler.gamma,
        patience=cfg.TRAIN.lr_scheduler.scheduler_patience,
        mode="min",
    )
elif cfg.TRAIN.lr_scheduler.type == "CosineAnnealingDecay":
    lr_scheduler = paddle.optimizer.lr.CosineAnnealingDecay(
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        T_max=cfg.TRAIN.lr_scheduler.scheduler_T_max,
    )()
elif cfg.TRAIN.lr_scheduler.type == "StepDecay":
    lr_scheduler = ppsci.optimizer.lr_scheduler.Step(
        epochs=cfg.TRAIN.lr_scheduler.epochs,
        iters_per_epoch=ITERS_PER_EPOCH,
        learning_rate=cfg.TRAIN.lr_scheduler.learning_rate,
        step_size=cfg.TRAIN.lr_scheduler.step_size,
        gamma=cfg.TRAIN.lr_scheduler.gamma,
    )()
else:
    raise ValueError(f"Got scheduler={cfg.TRAIN.lr_scheduler.type}")
optimizer = ppsci.optimizer.Adam(lr_scheduler, weight_decay=cfg.TRAIN.wd)(model)

5.2.4 评估器构建

本案例训练过程中会按照一定的训练轮数间隔，使用验证集评估当前模型的训练情况，需要使用 SupervisedValidator 构建评估器。代码如下：

examples/neuraloperator/train_sfno.py
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
# set eval dataloader config
eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "SphericalSWEDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "data_split": "test_32x64",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_64 = {
    "dataset": {
        "name": "SphericalSWEDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "data_split": "test_64x128",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

l2_eval_metric = metric.LpLoss(d=2, p=2, reduce_dims=[0, 1])
sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={"l2": ppsci.metric.FunctionalMetric(l2_eval_metric)},
    name="Sup_Validator_32x64",
)

sup_validator_64 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_64,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={"l2": ppsci.metric.FunctionalMetric(l2_eval_metric)},
    name="Sup_Validator_64x128",
)

validator = {
    sup_validator_32.name: sup_validator_32,
    sup_validator_64.name: sup_validator_64,
}

SupervisedValidator 评估器与 SupervisedConstraint 比较相似，不同的是评估器需要设置评价指标 metric，在这里使用了自定义的评价指标是 LpLoss。

5.2.5 模型训练与评估

完成上述设置之后，只需要将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动训练、评估。

examples/neuraloperator/train_sfno.py
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    lr_scheduler,
    cfg.TRAIN.epochs,
    ITERS_PER_EPOCH,
    eval_during_train=cfg.TRAIN.eval_during_train,
    seed=cfg.seed,
    validator=validator,
    compute_metric_by_batch=cfg.EVAL.compute_metric_by_batch,
    eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
    pretrained_model_path=cfg.TRAIN.pretrained_model_path,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

5.3 模型评估可视化

5.3.1 测试集上评估模型

构建模型的代码为：

examples/neuraloperator/train_sfno.py
217
218
219
model = ppsci.arch.SFNONet(
    **cfg.MODEL,
)

构建评估器的代码为：

examples/neuraloperator/train_sfno.py
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
# set eval dataloader config
eval_dataloader_cfg_32 = {
    "dataset": {
        "name": "SphericalSWEDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "data_split": "test_32x64",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

eval_dataloader_cfg_64 = {
    "dataset": {
        "name": "SphericalSWEDataset",
        "data_dir": cfg.FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.DATASET.label_keys,
        "train_resolution": cfg.DATASET.train_resolution,
        "test_resolutions": cfg.DATASET.test_resolutions,
        "data_split": "test_64x128",
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 0,
}

train_loss = metric.LpLoss_train(d=2, p=2, reduce_dims=[0, 1])

l2_eval_metric = metric.LpLoss(d=2, p=2, reduce_dims=[0, 1])
sup_validator_32 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_32,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={"l2": ppsci.metric.FunctionalMetric(l2_eval_metric)},
    name="Sup_Validator_32x64",
)

sup_validator_64 = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg_64,
    loss=ppsci.loss.FunctionalLoss(train_loss),
    metric={"l2": ppsci.metric.FunctionalMetric(l2_eval_metric)},
    name="Sup_Validator_64x128",
)

validator = {
    sup_validator_32.name: sup_validator_32,
    sup_validator_64.name: sup_validator_64,
}

5.3.2 模型导出

构建模型的代码为：

examples/neuraloperator/train_sfno.py
237
238
239
240
# set model
model = ppsci.arch.SFNONet(
    **cfg.MODEL,
)

实例化 ppsci.solver.Solver：

examples/neuraloperator/train_sfno.py
242
243
244
245
246
# initialize solver
solver = ppsci.solver.Solver(
    model,
    pretrained_model_path=cfg.INFER.pretrained_model_path,
)

构建模型输入格式并导出静态模型：

examples/neuraloperator/train_sfno.py
247
248
249
250
251
252
253
254
255
256
# export model
from paddle.static import InputSpec

input_spec = [
    {
        key: InputSpec([None, 3, 32, 64], "float32", name=key)
        for key in model.input_keys
    },
]
solver.export(input_spec, cfg.INFER.export_path)

InputSpec 函数中第一个设置模型输入尺寸，第二个参数设置输入数据类型，第三个设置输入数据的 Key.

5.3.3 模型推理

创建预测器:

examples/neuraloperator/train_sfno.py
261
262
263
import predictor

predictor = predictor.SFNOPredictor(cfg)

准备预测数据：

examples/neuraloperator/train_sfno.py
265
266
267
data = np.load(cfg.INFER.data_path, allow_pickle=True).item()
input_data = data["x"][0].reshape(1, *data["x"].shape[1:]).astype("float32")
label = data["y"][0][0, ...].astype("float32")

进行模型预测与预测值显示:

examples/neuraloperator/train_sfno.py
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
pred_data = predictor.predict(input_data, cfg.INFER.batch_size)

fig = plt.figure(figsize=(7, 7))
ax = fig.add_subplot(1, 3, 1)
ax.imshow(input_data.squeeze()[0, ...])
ax.set_title("k(x)")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 2)
ax.imshow(label)
ax.set_title("Ground-truth y")
plt.xticks([], [])
plt.yticks([], [])

ax = fig.add_subplot(1, 3, 3)
ax.imshow(pred_data.squeeze()[0, ...])
ax.set_title("Model prediction")
plt.xticks([], [])
plt.yticks([], [])
plt.savefig(cfg.output_dir)
logger.message("save success")
plt.close(fig)

6. 结果展示

下图展示了 TFNO 对 Darcy-flow 数据的预测结果和真值结果。 k(x) 的黑色区域就是可以渗透的地方，白色为不可渗透的区域。右侧是目标结果，颜色越亮，压力越大。

TFNO-predict

TFNO 的预测结果（"Model prediction"）与真值结果（"Ground-truth y"）

下图展示了 UNO 对 Darcy-flow 数据的预测结果和真值结果。

UNO-predict

UNO 的预测结果（"Model prediction"）与真值结果（"Ground-truth y"）

下图展示了 SFNO 对 SWE 数据的预测结果和真值结果。

SFNO-predict

SFNO的预测结果（"Model prediction"）与真值结果（"Ground-truth y"）