跳转至

Transolver

Note

运行本案例前请先安装相关依赖库: pip install -r requirements.txt

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip
# windows
# curl https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip -o mlcfd_data.zip
unzip mlcfd_data.zip
python main.py

Note

第一次运行时，会对 mlcfd_data 进行预处理，大约需要一小时，请耐心等待

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip
# windows
# curl https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip -o mlcfd_data.zip
unzip mlcfd_data.zip
python main.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/transolver/transolver_pretrained.pdparams

Note

第一次运行时，会对 mlcfd_data 进行预处理，大约需要一小时，请耐心等待

python main.py mode=export

# linux
wget -c https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip
# windows
# curl https://paddle-org.bj.bcebos.com/paddlecfd/datasets/pptransformer/mlcfd_data.zip -o mlcfd_data.zip
unzip mlcfd_data.zip
python main.py mode=infer

Note

第一次运行时，会对 mlcfd_data 进行预处理，大约需要一小时，请耐心等待

预训练模型 指标
transolver_pretrained.pdparams rho_d:, 0.99314
c_d: 0.01136
relative l2 error of press: 0.07829
relative l2 error of velocity: 0.02304
press: 4.95888
velocity: [0.12163974 0.14851639 0.41583335] 0.26443

1. 背景简介

Transolver 是一个基于 Transformer 架构的神经算子模型,用于学习偏微分方程(PDE)的解算子。该模型的核心创新在于其物理注意力机制(Physics Attention),能够高效地处理不规则网格上的物理场预测问题。

相比传统的 Transformer 模型,Transolver 具有以下特点:

  • 物理感知的注意力机制: 通过切片(slice)技术将不规则网格点聚合成规则的表示,既保留了空间物理信息,又大幅降低了计算复杂度
  • 灵活的几何适应性: 能够处理任意形状的几何体和非结构化网格
  • 高效的计算性能: 通过切片注意力机制,将计算复杂度从 \(O(N^2)\) 降低到 \(O(NG)\),其中 \(N\) 是网格点数,\(G\) 是切片数

本案例使用 Transolver 模型在 ShapeNet Car 数据集上学习汽车外流场的速度场和压力场分布,这是一个典型的计算流体动力学(CFD)代理建模问题。通过学习大量的汽车形状与对应的流场数据,模型可以快速预测新汽车形状的流场分布,从而大幅降低 CFD 仿真的计算成本。

2. 问题定义

本案例的目标是建立汽车几何形状到其周围流场(速度场和压力场)的映射关系。具体而言:

  • 输入: 汽车表面及周围空间的网格点坐标 \(\mathbf{x} \in \mathbb{R}^{N \times 7}\),其中 \(N\) 为网格点数量,7 个维度包括空间坐标、法向量等几何信息
  • 输出: 每个网格点处的速度矢量 \(\mathbf{v} \in \mathbb{R}^{N \times 3}\) 和压力标量 \(p \in \mathbb{R}^{N \times 1}\)

训练目标是最小化预测流场与真实 CFD 仿真结果之间的均方误差,同时确保模型能够准确预测汽车的阻力系数等关键空气动力学参数。

3. 问题求解

接下来开始讲解如何将问题一步一步地转化为 PaddleScience 代码,用深度学习的方法求解该问题。 为了快速理解 PaddleScience,接下来仅对模型构建、约束构建、优化器构建等关键步骤进行阐述,而其余细节请参考 API文档。

3.1 模型构建

在本问题中,需要建立从网格点坐标 \(\mathbf{x}\) 到流场变量 \((\mathbf{v}, p)\) 的映射函数 \(f: \mathbb{R}^{N \times 7} \to \mathbb{R}^{N \times 4}\),即:

\[ (\mathbf{v}, p) = f(\mathbf{x}) \]

这里使用 Transolver 模型来表示这个映射函数,用 PaddleScience 代码表示如下:

25
26
27
def train(cfg: DictConfig):
    # set model
    model = ppsci.arch.Transolver(**cfg.MODEL)

为了在计算时准确快速地访问具体变量的值,这里指定网络模型的输入变量名是 ["x"],输出变量名是 ["velo_vec", "press"],这些命名与后续代码保持一致。

模型的详细配置如下:

43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
MODEL:
  input_keys: [x]
  output_keys: [velo_vec, press] # vel: 3, press: 1
  space_dim: 7
  n_layers: 8
  n_hidden: 256
  dropout: 0
  n_head: 8
  act: gelu
  mlp_ratio: 2
  fun_dim: 0
  out_dim: [3, 1]
  slice_num: 32
  ref: 8
  unified_pos: false

其中:

  • space_dim: 输入空间维度,设为 7(包括空间坐标、法向量等几何信息)
  • n_layers: Transformer 层数,设为 8
  • n_hidden: 隐藏层维度,设为 256
  • n_head: 多头注意力的头数,设为 8
  • dropout: Dropout 比率,设为 0
  • act: 激活函数,使用 gelu
  • mlp_ratio: MLP 层的隐藏层扩展比例,设为 2
  • out_dim: 输出维度列表,[3, 1] 分别对应速度场(3维)和压力场(1维)
  • slice_num: 切片注意力机制中的切片数量,设为 32,用于降低计算复杂度
  • ref: 参考网格的分辨率,设为 8
  • unified_pos: 是否使用统一的位置编码,设为 False

3.2 模型架构详解

Transolver 的核心架构包括以下几个关键组件:

3.2.1 预处理层 (Preprocessing)

将输入的几何特征映射到隐藏空间:

\[ h = \text{MLP}(\text{concat}(f, x)) \]

其中 \(f\) 是函数特征(如初始场),\(x\) 是空间坐标。

3.2.2 物理注意力机制 (Physics Attention)

这是 Transolver 的核心创新,通过以下步骤实现高效的全局信息交互:

  1. 切片聚合: 将 \(N\) 个不规则网格点聚合成 \(G\) 个切片表示
\[ S = \text{softmax}\left(\frac{W_s(h)}{\tau}\right)^T h \]

其中 \(S \in \mathbb{R}^{G \times D}\) 是切片表示,\(\tau\) 是可学习的温度参数。

  1. 切片上的自注意力: 在切片表示上进行标准的 Transformer 注意力计算
\[ \text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V \]

其中 \(Q = W_q S, K = W_k S, V = W_v S\)。

  1. 反聚合: 将切片表示映射回原始网格点
\[ h' = S \cdot \text{Attn}(Q, K, V) \cdot W \]

通过这种机制,计算复杂度从 \(O(N^2)\) 降低到 \(O(NG + G^2)\),当 \(G \ll N\) 时可显著提升效率。

3.2.3 前馈网络 (MLP)

每个 Transformer 块后接一个前馈网络:

\[ \text{FFN}(h) = W_2 \cdot \text{GELU}(W_1 h + b_1) + b_2 \]

3.2.4 层归一化与残差连接

每个子层都使用层归一化(LayerNorm)和残差连接:

\[ \begin{aligned} h^{(l+1)} &= \text{Attn}(\text{LN}(h^{(l)})) + h^{(l)} \\ h^{(l+1)} &= \text{FFN}(\text{LN}(h^{(l+1)})) + h^{(l+1)} \end{aligned} \]

3.3 数据加载

本案例使用 ShapeNet Car 数据集,包含不同汽车形状的 CFD 仿真结果。数据加载代码如下:

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# set constraint
train_data, val_data, coef_norm = load_train_val_fold(
    cfg.DATA.data_dir,
    cfg.DATA.val_fold_id,
    cfg.DATA.save_dir,
    preprocessed=cfg.DATA.preprocessed,
)
train_dataset = ShapeNetCarDataset(
    cfg.DATA.input_keys,
    cfg.DATA.label_keys,
    train_data,
    use_cfd_mesh=cfg.DATA.use_cfd_mesh,
    r=cfg.DATA.r,
    training=True,
)

数据集配置:

31
32
33
34
35
36
37
38
39
DATA:
  input_keys: [x]
  label_keys: [velo_vec, press]
  data_dir: ./mlcfd_data/training_data
  save_dir: ./mlcfd_data/preprocessed_data
  val_fold_id: 0
  preprocessed: false
  use_cfd_mesh: false
  r: 0.2

其中:

  • data_dir: 原始数据目录
  • save_dir: 预处理后数据保存目录
  • val_fold_id: 用于交叉验证的折数 ID
  • preprocessed: 是否使用预处理数据
  • r: 数据下采样比例

3.4 约束构建

本案例使用监督学习约束,通过最小化预测流场与真实流场的误差来训练模型:

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
def loss_func(output, label, _):
    velo_label = label[cfg.DATA.label_keys[0]]
    press_label = label[cfg.DATA.label_keys[1]]
    surf_mask = label["surf"]
    mse_velo = F.mse_loss(output[cfg.DATA.label_keys[0]], velo_label)
    mse_press = F.mse_loss(
        output[cfg.MODEL.output_keys[1]][surf_mask], press_label[surf_mask]
    )

    return {
        "mse_velo_press": mse_velo + cfg.TRAIN.press_weight * mse_press,
    }

sup_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": train_dataset,
        "batch_size": cfg.TRAIN.batch_size,
        "num_workers": 0,
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": True,
        },
    },
    loss=ppsci.loss.FunctionalLoss(loss_func),
    name="Sup",
)

# wrap constraints together
constraint = {
    sup_constraint.name: sup_constraint,
}

损失函数包含两部分:

  1. 速度场的均方误差: \(\mathcal{L}_{velo} = \text{MSE}(\mathbf{v}_{pred}, \mathbf{v}_{true})\)
  2. 表面压力场的均方误差: \(\mathcal{L}_{press} = \text{MSE}(p_{pred}|_{surf}, p_{true}|_{surf})\)

总损失为: \(\mathcal{L} = \mathcal{L}_{velo} + w_{press} \cdot \mathcal{L}_{press}\),其中 \(w_{press}\) 为压力损失权重。

3.5 优化器构建

使用 Adam 优化器配合指数衰减学习率策略:

78
79
80
81
82
83
84
85
86
87
88
89
# set optimizer
## Slightly differnt from transolver's lr setting(OneCycleLR)
lr = ppsci.optimizer.lr_scheduler.ExponentialDecay(
    cfg.TRAIN.epochs,
    len(sup_constraint.data_loader),
    cfg.TRAIN.lr.max_lr,
    gamma=cfg.TRAIN.lr.gamma,
    decay_steps=cfg.TRAIN.lr.decay_steps,
    warmup_epoch=cfg.TRAIN.lr.warmup_epoch,
    warmup_start_lr=cfg.TRAIN.lr.max_lr / 25.0,
)()
optimizer = ppsci.optimizer.Adam(lr)(model)

学习率配置:

63
64
65
66
67
lr:
  warmup_epoch: 60
  max_lr: 1e-3
  gamma: 0.999973
  decay_steps: 1

其中包含:

  • 预热阶段(warmup_epoch): 前 60 个 epoch 学习率从 \(\frac{lr_{max}}{25}\) 逐渐增加到 \(lr_{max}\)
  • 衰减阶段: 之后每个 step 学习率按 \(lr = lr_{max} \times \gamma^{step}\) 衰减

3.6 评估器构建

在训练过程中使用验证集评估模型性能:

 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 validator
val_dataset = ShapeNetCarDataset(
    cfg.DATA.input_keys,
    cfg.DATA.label_keys,
    val_data,
    use_cfd_mesh=cfg.DATA.use_cfd_mesh,
    r=cfg.DATA.r,
    training=False,
)

def val_metric_func(output, label):
    velo_label = label[cfg.DATA.label_keys[0]]
    press_label = label[cfg.DATA.label_keys[1]]
    surf_mask = label["surf"]
    loss_velo_vec = F.mse_loss(
        output[cfg.DATA.label_keys[0]], velo_label, "none"
    ).mean(axis=0)
    loss_velo = loss_velo_vec.mean()
    loss_press = F.mse_loss(
        output[cfg.DATA.label_keys[1]][surf_mask], press_label[surf_mask]
    )

    return {
        "press": loss_press,
        "velo_vec": loss_velo,
    }

validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": val_dataset,
        "batch_size": cfg.EVAL.batch_size,
        "num_workers": 0,
    },
    metric={
        "mse": ppsci.metric.FunctionalMetric(val_metric_func),
    },
    name="validator",
)
validator = {validator.name: validator}

评估指标包括:

  • 速度场的均方误差
  • 表面压力场的均方误差

3.7 模型训练与评估

完成上述设置后,将实例化的对象传递给 ppsci.solver.Solver,然后启动训练和评估:

140
141
142
solver.eval()
# train model
solver.train()

3.8 模型评估

评估阶段除了计算常规的误差指标外,还会计算阻力系数的预测误差和斯皮尔曼相关系数:

145
146
147
148
149
150
151
152
153
154
155
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
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
def eval_on_dataloadr(model_forward, dataloader, coef_norm, val_list):
    with paddle.no_grad():
        l2errs_press = []
        l2errs_velo = []
        mses_press = []
        mses_velo_var = []
        gt_coef_list = []
        pred_coef_list = []
        coef_error = 0
        index = 0
        pbar = tqdm(dataloader, desc="Testing", unit="batch")
        for i, (inp, label, _) in enumerate(pbar, start=1):
            out = model_forward(inp)
            velo_vec = out["velo_vec"]
            press = out["press"]
            targets_velo_vec = label["velo_vec"]
            targets_press = label["press"]
            surf_mask = label["surf"]

            if coef_norm is not None:
                mean = paddle.tensor(coef_norm[2], dtype=dtype)
                std = paddle.tensor(coef_norm[3], dtype=dtype)

                pred_press: paddle.Tensor = press[surf_mask] * std[-1] + mean[-1]
                gt_press: paddle.Tensor = targets_press[surf_mask] * std[-1] + mean[-1]

                pred_surf_velo: paddle.Tensor = (
                    velo_vec[surf_mask] * std[:-1] + mean[:-1]
                )
                gt_surf_velo: paddle.Tensor = (
                    targets_velo_vec[surf_mask] * std[:-1] + mean[:-1]
                )

                pred_velo: paddle.Tensor = velo_vec[~surf_mask] * std[:-1] + mean[:-1]
                gt_velo: paddle.Tensor = (
                    targets_velo_vec[~surf_mask] * std[:-1] + mean[:-1]
                )

                # out_denorm: paddle.Tensor = out * std + mean
                # y_denorm: paddle.Tensor = targets * std + mean
                # np.save('./results/' + args.cfd_model + '/' + str(index) + '_pred.npy', out_denorm.numpy())
                # np.save('./results/' + args.cfd_model + '/' + str(index) + '_gt.npy', y_denorm.numpy())

            pred_coef = cal_coefficient(
                val_list[index].split("/")[1],
                pred_press[:, None].numpy(),
                pred_surf_velo.numpy(),
            )
            gt_coef = cal_coefficient(
                val_list[index].split("/")[1],
                gt_press[:, None].numpy(),
                gt_surf_velo.numpy(),
            )

            gt_coef_list.append(gt_coef)
            pred_coef_list.append(pred_coef)
            coef_error += abs(pred_coef - gt_coef) / gt_coef
            pbar.set_postfix(
                {
                    "batch": f"{i}/{len(dataloader)}, coef_error: {coef_error / (index + 1):.10f}",
                }
            )

            l2err_press = paddle.norm(pred_press - gt_press) / paddle.norm(gt_press)
            l2err_velo = paddle.norm(pred_velo - gt_velo) / paddle.norm(gt_velo)

            mse_press = F.mse_loss(
                press[surf_mask], targets_press[surf_mask], "none"
            ).mean(axis=0)
            mse_velo_var = F.mse_loss(
                velo_vec[~surf_mask], targets_velo_vec[~surf_mask], "none"
            ).mean(axis=0)

            l2errs_press.append(l2err_press.numpy())
            l2errs_velo.append(l2err_velo.numpy())
            mses_press.append(mse_press.numpy())
            mses_velo_var.append(mse_velo_var.numpy())
            index += 1

        gt_coef_list = np.array(gt_coef_list)
        pred_coef_list = np.array(pred_coef_list)
        spear = sc.stats.spearmanr(gt_coef_list, pred_coef_list)[0]
        logger.info(f"rho_d:, {spear:.5f}")
        logger.info(f"c_d: {coef_error / index:.5f}")
        l2err_press = np.mean(l2errs_press)
        l2err_velo = np.mean(l2errs_velo)
        rmse_press = np.sqrt(np.mean(mses_press))
        rmse_velo_var = np.sqrt(np.mean(mses_velo_var, axis=0))
        if coef_norm is not None:
            rmse_press *= coef_norm[3][-1]
            rmse_velo_var *= coef_norm[3][:-1]
        logger.info(f"relative l2 error of press: {l2err_press:.5f}")
        logger.info(f"relative l2 error of velocity: {l2err_velo:.5f}")
        logger.info(f"press: {rmse_press:.5f}")
        logger.info(
            f"velocity: {rmse_velo_var} {np.sqrt(np.mean(np.square(rmse_velo_var))):.5f}"
        )

评估指标包括:

  • 相对 L2 误差: 衡量预测场与真实场的整体偏差
  • 均方根误差 (RMSE): 评估逐点预测精度
  • 阻力系数误差: 评估关键空气动力学参数的预测精度
  • 斯皮尔曼相关系数 (\(\rho_d\)): 评估阻力系数排序的相关性

4. 完整代码

main.py
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 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
 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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
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
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
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
"""
Reference: https://github.com/thuml/Transolver
"""
from __future__ import annotations

import hydra
import numpy as np
import paddle
import paddle.nn.functional as F
import scipy as sc
from drag_coefficient import cal_coefficient
from omegaconf import DictConfig
from shapenet_car import ShapeNetCarDataset
from shapenet_car import load_train_val_fold
from shapenet_car import load_train_val_fold_file
from tqdm import tqdm

import ppsci
from ppsci.utils import logger
from ppsci.utils import save_load

dtype = paddle.get_default_dtype()


def train(cfg: DictConfig):
    # set model
    model = ppsci.arch.Transolver(**cfg.MODEL)

    # set constraint
    train_data, val_data, coef_norm = load_train_val_fold(
        cfg.DATA.data_dir,
        cfg.DATA.val_fold_id,
        cfg.DATA.save_dir,
        preprocessed=cfg.DATA.preprocessed,
    )
    train_dataset = ShapeNetCarDataset(
        cfg.DATA.input_keys,
        cfg.DATA.label_keys,
        train_data,
        use_cfd_mesh=cfg.DATA.use_cfd_mesh,
        r=cfg.DATA.r,
        training=True,
    )

    def loss_func(output, label, _):
        velo_label = label[cfg.DATA.label_keys[0]]
        press_label = label[cfg.DATA.label_keys[1]]
        surf_mask = label["surf"]
        mse_velo = F.mse_loss(output[cfg.DATA.label_keys[0]], velo_label)
        mse_press = F.mse_loss(
            output[cfg.MODEL.output_keys[1]][surf_mask], press_label[surf_mask]
        )

        return {
            "mse_velo_press": mse_velo + cfg.TRAIN.press_weight * mse_press,
        }

    sup_constraint = ppsci.constraint.SupervisedConstraint(
        {
            "dataset": train_dataset,
            "batch_size": cfg.TRAIN.batch_size,
            "num_workers": 0,
            "sampler": {
                "name": "BatchSampler",
                "drop_last": True,
                "shuffle": True,
            },
        },
        loss=ppsci.loss.FunctionalLoss(loss_func),
        name="Sup",
    )

    # wrap constraints together
    constraint = {
        sup_constraint.name: sup_constraint,
    }

    # set optimizer
    ## Slightly differnt from transolver's lr setting(OneCycleLR)
    lr = ppsci.optimizer.lr_scheduler.ExponentialDecay(
        cfg.TRAIN.epochs,
        len(sup_constraint.data_loader),
        cfg.TRAIN.lr.max_lr,
        gamma=cfg.TRAIN.lr.gamma,
        decay_steps=cfg.TRAIN.lr.decay_steps,
        warmup_epoch=cfg.TRAIN.lr.warmup_epoch,
        warmup_start_lr=cfg.TRAIN.lr.max_lr / 25.0,
    )()
    optimizer = ppsci.optimizer.Adam(lr)(model)

    # set validator
    val_dataset = ShapeNetCarDataset(
        cfg.DATA.input_keys,
        cfg.DATA.label_keys,
        val_data,
        use_cfd_mesh=cfg.DATA.use_cfd_mesh,
        r=cfg.DATA.r,
        training=False,
    )

    def val_metric_func(output, label):
        velo_label = label[cfg.DATA.label_keys[0]]
        press_label = label[cfg.DATA.label_keys[1]]
        surf_mask = label["surf"]
        loss_velo_vec = F.mse_loss(
            output[cfg.DATA.label_keys[0]], velo_label, "none"
        ).mean(axis=0)
        loss_velo = loss_velo_vec.mean()
        loss_press = F.mse_loss(
            output[cfg.DATA.label_keys[1]][surf_mask], press_label[surf_mask]
        )

        return {
            "press": loss_press,
            "velo_vec": loss_velo,
        }

    validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": val_dataset,
            "batch_size": cfg.EVAL.batch_size,
            "num_workers": 0,
        },
        metric={
            "mse": ppsci.metric.FunctionalMetric(val_metric_func),
        },
        name="validator",
    )
    validator = {validator.name: validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        optimizer=optimizer,
        validator=validator,
        cfg=cfg,
    )

    solver.eval()
    # train model
    solver.train()


def eval_on_dataloadr(model_forward, dataloader, coef_norm, val_list):
    with paddle.no_grad():
        l2errs_press = []
        l2errs_velo = []
        mses_press = []
        mses_velo_var = []
        gt_coef_list = []
        pred_coef_list = []
        coef_error = 0
        index = 0
        pbar = tqdm(dataloader, desc="Testing", unit="batch")
        for i, (inp, label, _) in enumerate(pbar, start=1):
            out = model_forward(inp)
            velo_vec = out["velo_vec"]
            press = out["press"]
            targets_velo_vec = label["velo_vec"]
            targets_press = label["press"]
            surf_mask = label["surf"]

            if coef_norm is not None:
                mean = paddle.tensor(coef_norm[2], dtype=dtype)
                std = paddle.tensor(coef_norm[3], dtype=dtype)

                pred_press: paddle.Tensor = press[surf_mask] * std[-1] + mean[-1]
                gt_press: paddle.Tensor = targets_press[surf_mask] * std[-1] + mean[-1]

                pred_surf_velo: paddle.Tensor = (
                    velo_vec[surf_mask] * std[:-1] + mean[:-1]
                )
                gt_surf_velo: paddle.Tensor = (
                    targets_velo_vec[surf_mask] * std[:-1] + mean[:-1]
                )

                pred_velo: paddle.Tensor = velo_vec[~surf_mask] * std[:-1] + mean[:-1]
                gt_velo: paddle.Tensor = (
                    targets_velo_vec[~surf_mask] * std[:-1] + mean[:-1]
                )

                # out_denorm: paddle.Tensor = out * std + mean
                # y_denorm: paddle.Tensor = targets * std + mean
                # np.save('./results/' + args.cfd_model + '/' + str(index) + '_pred.npy', out_denorm.numpy())
                # np.save('./results/' + args.cfd_model + '/' + str(index) + '_gt.npy', y_denorm.numpy())

            pred_coef = cal_coefficient(
                val_list[index].split("/")[1],
                pred_press[:, None].numpy(),
                pred_surf_velo.numpy(),
            )
            gt_coef = cal_coefficient(
                val_list[index].split("/")[1],
                gt_press[:, None].numpy(),
                gt_surf_velo.numpy(),
            )

            gt_coef_list.append(gt_coef)
            pred_coef_list.append(pred_coef)
            coef_error += abs(pred_coef - gt_coef) / gt_coef
            pbar.set_postfix(
                {
                    "batch": f"{i}/{len(dataloader)}, coef_error: {coef_error / (index + 1):.10f}",
                }
            )

            l2err_press = paddle.norm(pred_press - gt_press) / paddle.norm(gt_press)
            l2err_velo = paddle.norm(pred_velo - gt_velo) / paddle.norm(gt_velo)

            mse_press = F.mse_loss(
                press[surf_mask], targets_press[surf_mask], "none"
            ).mean(axis=0)
            mse_velo_var = F.mse_loss(
                velo_vec[~surf_mask], targets_velo_vec[~surf_mask], "none"
            ).mean(axis=0)

            l2errs_press.append(l2err_press.numpy())
            l2errs_velo.append(l2err_velo.numpy())
            mses_press.append(mse_press.numpy())
            mses_velo_var.append(mse_velo_var.numpy())
            index += 1

        gt_coef_list = np.array(gt_coef_list)
        pred_coef_list = np.array(pred_coef_list)
        spear = sc.stats.spearmanr(gt_coef_list, pred_coef_list)[0]
        logger.info(f"rho_d:, {spear:.5f}")
        logger.info(f"c_d: {coef_error / index:.5f}")
        l2err_press = np.mean(l2errs_press)
        l2err_velo = np.mean(l2errs_velo)
        rmse_press = np.sqrt(np.mean(mses_press))
        rmse_velo_var = np.sqrt(np.mean(mses_velo_var, axis=0))
        if coef_norm is not None:
            rmse_press *= coef_norm[3][-1]
            rmse_velo_var *= coef_norm[3][:-1]
        logger.info(f"relative l2 error of press: {l2err_press:.5f}")
        logger.info(f"relative l2 error of velocity: {l2err_velo:.5f}")
        logger.info(f"press: {rmse_press:.5f}")
        logger.info(
            f"velocity: {rmse_velo_var} {np.sqrt(np.mean(np.square(rmse_velo_var))):.5f}"
        )


def evaluate(cfg: DictConfig):
    # set model
    model = ppsci.arch.Transolver(**cfg.MODEL)

    # load pretrained model
    save_load.load_pretrain(model, cfg.EVAL.pretrained_model_path)
    model.eval()

    # evaluate manually
    _, val_data, coef_norm, val_list = load_train_val_fold_file(
        cfg.DATA.data_dir,
        cfg.DATA.val_fold_id,
        cfg.DATA.save_dir,
        preprocessed=cfg.DATA.preprocessed,
    )
    test_dataset = ShapeNetCarDataset(
        cfg.DATA.input_keys,
        cfg.DATA.label_keys,
        val_data,
        use_cfd_mesh=cfg.DATA.use_cfd_mesh,
        r=cfg.DATA.r,
        training=False,
    )
    test_dataloader = ppsci.data.build_dataloader(
        test_dataset,
        {
            "batch_size": cfg.EVAL.batch_size,
            "num_workers": 0,
        },
    )

    eval_on_dataloadr(model, test_dataloader, coef_norm, val_list)


def export(cfg: DictConfig):
    # set model
    model = ppsci.arch.Transolver(**cfg.MODEL)

    # initialize solver
    solver = ppsci.solver.Solver(model, cfg=cfg)
    # export model
    from paddle.static import InputSpec

    input_spec = [
        {
            model.input_keys[0]: InputSpec(
                [1, None, 7], "float32", name=model.input_keys[0]
            ),
        },
    ]
    import einops

    solver.export(
        input_spec, cfg.INFER.export_path, with_onnx=False, ignore_modules=[einops]
    )


def inference(cfg: DictConfig):
    from deploy import python_infer

    predictor = python_infer.GeneralPredictor(cfg)

    # inference manually
    _, val_data, coef_norm, val_list = load_train_val_fold_file(
        cfg.DATA.data_dir,
        cfg.DATA.val_fold_id,
        cfg.DATA.save_dir,
        preprocessed=cfg.DATA.preprocessed,
    )
    test_dataset = ShapeNetCarDataset(
        cfg.DATA.input_keys,
        cfg.DATA.label_keys,
        val_data,
        use_cfd_mesh=cfg.DATA.use_cfd_mesh,
        r=cfg.DATA.r,
        training=False,
    )
    test_dataloader = ppsci.data.build_dataloader(
        test_dataset,
        {
            "batch_size": cfg.TRAIN.batch_size,
            "num_workers": 0,
        },
    )

    def wrap_predict(x):
        raw_out = predictor.predict(
            {k: v.numpy() for k, v in x.items()}, batch_size=None
        )
        return {
            store_key: paddle.tensor(raw_out[infer_key])
            for store_key, infer_key in zip(cfg.MODEL.output_keys[::-1], raw_out.keys())
        }

    eval_on_dataloadr(wrap_predict, test_dataloader, coef_norm, val_list)


@hydra.main(version_base=None, config_path="./conf", config_name="shapenet_car.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "export":
        export(cfg)
    elif cfg.mode == "infer":
        inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

5. 结果展示

模型训练完成后,可以在验证集上评估预测性能。主要评估指标包括:

  • 速度场相对 L2 误差: 衡量速度场预测的整体精度
  • 压力场相对 L2 误差: 衡量压力场预测的整体精度
  • 阻力系数相对误差: 评估关键气动参数的预测准确性
  • 斯皮尔曼相关系数: 评估不同汽车形状阻力系数排序的准确性

通过训练,Transolver 模型能够以较高精度预测汽车周围的流场分布,相比传统 CFD 仿真可以大幅降低计算时间,为汽车气动外形优化提供快速的代理模型。

6. 参考资料