跳转至

Lorenz System

AI Studio快速体验

# linux
wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5 -P ./datasets/
wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5 -P ./datasets/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5 --output ./datasets/lorenz_training_rk.hdf5
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5 --output ./datasets/lorenz_valid_rk.hdf5
python train_enn.py
python train_transformer.py

# linux
wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5 -P ./datasets/
wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5 -P ./datasets/
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_training_rk.hdf5 --output ./datasets/lorenz_training_rk.hdf5
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/transformer_physx/lorenz_valid_rk.hdf5 --output ./datasets/lorenz_valid_rk.hdf5
python train_enn.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/lorenz/lorenz_pretrained.pdparams
python train_transformer.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/lorenz/lorenz_transformer_pretrained.pdparams EMBEDDING_MODEL_PATH=https://paddle-org.bj.bcebos.com/paddlescience/models/lorenz/lorenz_pretrained.pdparams
模型 MSE
lorenz_transformer_pretrained.pdparams 0.054

1. 背景简介

Lorenz System，中文名称可译作“洛伦兹系统”，又称“洛伦兹混沌系统”，最早由美国气象学家爱德华·洛伦兹（Edward N.Lorenz）在1963年的一篇文章中提出。著名的“蝴蝶效应”，即“一只南美洲亚马逊河流域热带雨林中的蝴蝶，偶尔扇动几下翅膀，可以在两周以后引起美国得克萨斯州的一场龙卷风”，也是最早起源于这篇文章。洛伦兹系统的特点是在一定参数条件下展现出复杂、不确定的动态行为，包括对初始条件的敏感性和长期行为的不可预测性。这种混沌行为在自然界和许多实际应用领域中都存在，例如气候变化、股票市场波动等。洛伦兹系统对数值扰动极为敏感，是评估机器学习（深度学习）模型准确性的良好基准。

2. 问题定义

洛伦兹系统的状态方程：

\[ \begin{cases} \dfrac{\partial x}{\partial t} = \sigma(y - x), & \\ \dfrac{\partial y}{\partial t} = x(\rho - z) - y, & \\ \dfrac{\partial z}{\partial t} = xy - \beta z \end{cases} \]

当参数取以下值时，系统表现出经典的混沌特性：

\[\rho = 28, \sigma = 10, \beta = \frac{8}{3}\]

在这个案例中，要求给定初始时刻点的坐标，预测未来一段时间内点的运动轨迹。

3. 问题求解

接下来开始讲解如何基于 PaddleScience 代码，用深度学习的方法求解该问题。本案例基于论文 Transformers for Modeling Physical Systems 方法进行求解，接下来首先会对该论文的理论方法进行简单介绍，然后对使用的数据集进行介绍，最后对该方法两个训练步骤（Embedding 模型训练、Transformer 模型训练）的监督约束构建、模型构建等进行阐述，而其余细节请参考 API文档。

3.1 方法介绍

Transformer 结构在 NLP、CV 领域中取得了巨大的成功，但是其在建模物理系统方面还没有得到更多的探索。在 Transformers for Modeling Physical Systems 这篇文章中，作者提出了基于 Transformer 的网络结构用于建模物理系统。实验结果表明，提出的方法能够准确的建模不同的动态系统，并且比其他传统的方法更好。

如下图所示，该方法主要包含两个网络模型：Embedding 模型和 Transformer 模型。其中，Embedding 模型的 Encoder 模块负责将物理状态变量进行编码映射为编码向量，Decoder 模块则负责将编码向量映射为物理状态变量；Transformer 模型作用于编码空间，其输入是 Embedding 模型 Encoder 模块的输出，利用当前时刻的编码向量预测下一时刻的编码向量，预测得到的编码向量可以被 Embedding 模型的 Decoder 模块解码，得到对应的物理状态变量。在模型训练时，首先训练 Embedding 模型，然后将 Embedding 模型的参数冻结训练 Transformer 模型。关于该方法的细节请参考论文 Transformers for Modeling Physical Systems。

trphysx-arch

左：Embedding 网络结构，右：Transformer 网络结构

3.2 数据集介绍

数据集采用了 Transformer-Physx 中提供的数据。该数据集使用龙格－库塔（Runge-Kutta）传统数值求解方法得到，每个时间步大小为0.01，初始位置从以下范围中随机选取：

\[x_{0} \sim(-20, 20), y_{0} \sim(-20, 20), z_{0} \sim(10, 40)\]

数据集的划分如下：

数据集 时间序列的数量 时间步的数量 下载地址
训练集 2048 256 lorenz_training_rk.hdf5
验证集 64 1024 lorenz_valid_rk.hdf5

数据集官网为：https://zenodo.org/record/5148524#.ZDe77-xByrc

3.3 Embedding 模型

首先展示代码中定义的各个参数变量，每个参数的具体含义会在下面使用到时进行解释。

examples/conf/enn.yaml
26
27
28
29
30
31
32
33
34
TRAIN_BLOCK_SIZE: 16
VALID_BLOCK_SIZE: 32
TRAIN_FILE_PATH: ./datasets/lorenz_training_rk.hdf5
VALID_FILE_PATH: ./datasets/lorenz_valid_rk.hdf5

# model settings
MODEL:
  input_keys: ["states"]
  output_keys: ["pred_states", "recover_states"]

3.3.1 约束构建

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

examples/lorenz/train_enn.py
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
train_dataloader_cfg = {
    "dataset": {
        "name": "LorenzDataset",
        "file_path": cfg.TRAIN_FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.output_keys,
        "block_size": cfg.TRAIN_BLOCK_SIZE,
        "stride": 16,
        "weight_dict": {
            key: value for key, value in zip(cfg.MODEL.output_keys, weights)
        },
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": True,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
    "num_workers": 4,
}

其中，"dataset" 字段定义了使用的 Dataset 类名为 LorenzDataset，另外还指定了该类初始化时参数的取值：

  1. file_path：代表训练数据集的文件路径，指定为变量 train_file_path 的值；
  2. input_keys：代表模型输入数据的变量名称，此处填入变量 input_keys；
  3. label_keys：代表真实标签的变量名称，此处填入变量 output_keys；
  4. block_size：代表使用多长的时间步进行训练，指定为变量 train_block_size 的值；
  5. stride：代表连续的两个训练样本之间的时间步间隔，指定为16；
  6. weight_dict：代表模型输出各个变量与真实标签损失函数的权重，此处使用 output_keys、weights 生成。

"sampler" 字段定义了使用的 Sampler 类名为 BatchSampler，另外还指定了该类初始化时参数 drop_last、shuffle 均为 True。

train_dataloader_cfg 还定义了 batch_size、num_workers 的值。

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

examples/lorenz/train_enn.py
72
73
74
75
76
77
78
79
80
81
82
83
84
85
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    ppsci.loss.MSELossWithL2Decay(
        regularization_dict={
            regularization_key: 1.0e-1 * (cfg.TRAIN_BLOCK_SIZE - 1)
        }
    ),
    {
        key: lambda out, k=key: out[k]
        for key in cfg.MODEL.output_keys + (regularization_key,)
    },
    name="Sup",
)
constraint = {sup_constraint.name: sup_constraint}

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

第二个参数是损失函数的定义，这里使用带有 L2Decay 的 MSELoss，类名为 MSELossWithL2Decay，regularization_dict 设置了正则化的变量名称和对应的权重；

第三个参数表示在训练时如何计算需要被约束的中间变量，此处我们约束的变量就是网络的输出；

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

3.3.2 模型构建

在该案例中，Embedding 模型的输入输出都是物理空间中点的位置坐标 \((x, y, z)\) ，使用了全连接层实现 Embedding 模型，如下图所示。

lorenz_embedding

Embedding 网络模型

用 PaddleScience 代码表示如下：

examples/lorenz/train_enn.py
91
92
93
94
95
96
97
data_mean, data_std = get_mean_std(sup_constraint.data_loader.dataset.data)
model = ppsci.arch.LorenzEmbedding(
    cfg.MODEL.input_keys,
    cfg.MODEL.output_keys + (regularization_key,),
    data_mean,
    data_std,
)

其中，LorenzEmbedding 的前两个参数在前文中已有描述，这里不再赘述，网络模型的第三、四个参数是训练数据集的均值和方差，用于归一化输入数据。计算均值、方差的的代码表示如下：

examples/lorenz/train_enn.py
32
33
34
35
36
37
38
39
def get_mean_std(data: np.ndarray):
    mean = np.asarray(
        [np.mean(data[:, :, 0]), np.mean(data[:, :, 1]), np.mean(data[:, :, 2])]
    ).reshape(1, 3)
    std = np.asarray(
        [np.std(data[:, :, 0]), np.std(data[:, :, 1]), np.std(data[:, :, 2])]
    ).reshape(1, 3)
    return mean, std

3.3.3 学习率与优化器构建

本案例中使用的学习率方法为 ExponentialDecay ，学习率大小设置为0.001。优化器使用 Adam，梯度裁剪使用了 Paddle 内置的 ClipGradByGlobalNorm 方法。用 PaddleScience 代码表示如下

examples/lorenz/train_enn.py
 99
100
101
102
103
104
105
106
107
108
# init optimizer and lr scheduler
clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=0.1)
lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(
    iters_per_epoch=ITERS_PER_EPOCH,
    decay_steps=ITERS_PER_EPOCH,
    **cfg.TRAIN.lr_scheduler,
)()
optimizer = ppsci.optimizer.Adam(
    lr_scheduler, grad_clip=clip, **cfg.TRAIN.optimizer
)(model)

3.3.4 评估器构建

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

examples/lorenz/train_enn.py
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
eval_dataloader_cfg = {
    "dataset": {
        "name": "LorenzDataset",
        "file_path": cfg.VALID_FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.output_keys,
        "block_size": cfg.VALID_BLOCK_SIZE,
        "stride": 32,
        "weight_dict": {
            key: value for key, value in zip(cfg.MODEL.output_keys, weights)
        },
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 4,
}

mse_validator = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg,
    ppsci.loss.MSELoss(),
    metric={"MSE": ppsci.metric.MSE()},
    name="MSE_Validator",
)
validator = {mse_validator.name: mse_validator}

SupervisedValidator 评估器与 SupervisedConstraint 比较相似，不同的是评估器需要设置评价指标 metric，在这里使用 ppsci.metric.MSE 。

3.3.5 模型训练与评估

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

examples/lorenz/train_enn.py
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    lr_scheduler,
    cfg.TRAIN.epochs,
    ITERS_PER_EPOCH,
    eval_during_train=True,
    validator=validator,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

3.4 Transformer 模型

上文介绍了如何构建 Embedding 模型的训练、评估，在本节中将介绍如何使用训练好的 Embedding 模型训练 Transformer 模型。因为训练 Transformer 模型的步骤与训练 Embedding 模型的步骤基本相似，因此本节在两者的重复部分的各个参数不再详细介绍。首先将代码中定义的各个参数变量展示如下，每个参数的具体含义会在下面使用到时进行解释。

examples/lorenz/conf/transformer.yaml
36
37
38
39
40
41
42
43
# model settings
MODEL:
  input_keys: ["embeds"]
  output_keys: ["pred_embeds"]
  num_layers: 4
  num_ctx: 64
  embed_size: 32
  num_heads: 4

3.4.1 约束构建

Transformer 模型同样基于数据驱动的方法求解问题，因此需要使用 PaddleScience 内置的 SupervisedConstraint 构建监督约束。在定义约束之前，需要首先指定监督约束中用于数据加载的各个参数，代码如下：

examples/lorenz/train_transformer.py
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
train_dataloader_cfg = {
    "dataset": {
        "name": "LorenzDataset",
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.output_keys,
        "file_path": cfg.TRAIN_FILE_PATH,
        "block_size": cfg.TRAIN_BLOCK_SIZE,
        "stride": 64,
        "embedding_model": embedding_model,
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": True,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
    "num_workers": 4,
}

数据加载的各个参数与 Embedding 模型中的基本一致，不再赘述。需要说明的是由于 Transformer 模型训练的输入数据是 Embedding 模型 Encoder 模块的输出数据，因此我们将训练好的 Embedding 模型作为 LorenzDataset 的一个参数，在初始化时首先将训练数据映射到编码空间。

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

examples/lorenz/train_transformer.py
87
88
89
90
91
92
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    ppsci.loss.MSELoss(),
    name="Sup",
)
constraint = {sup_constraint.name: sup_constraint}

3.4.2 模型构建

在该案例中，Transformer 模型的输入输出都是编码空间中的向量，使用的 Transformer 结构如下：

lorenz_transformer

Transformer 网络模型

用 PaddleScience 代码表示如下：

examples/lorenz/train_transformer.py
98
model = ppsci.arch.PhysformerGPT2(**cfg.MODEL)

类 PhysformerGPT2 除了需要填入 input_keys、output_keys 外，还需要设置 Transformer 模型的层数 num_layers、上下文的大小 num_ctx、输入的 Embedding 向量的长度 embed_size、多头注意力机制的参数 num_heads，在这里填入的数值为4、64、32、4。

3.4.3 学习率与优化器构建

本案例中使用的学习率方法为 CosineWarmRestarts，学习率大小设置为0.001。优化器使用 Adam，梯度裁剪使用了 Paddle 内置的 ClipGradByGlobalNorm 方法。用 PaddleScience 代码表示如下：

examples/lorenz/train_transformer.py
101
102
103
104
105
106
107
clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=0.1)
lr_scheduler = ppsci.optimizer.lr_scheduler.CosineWarmRestarts(
    iters_per_epoch=ITERS_PER_EPOCH, **cfg.TRAIN.lr_scheduler
)()
optimizer = ppsci.optimizer.Adam(
    lr_scheduler, grad_clip=clip, **cfg.TRAIN.optimizer
)(model)

3.4.4 评估器构建

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

examples/lorenz/train_transformer.py
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
eval_dataloader_cfg = {
    "dataset": {
        "name": "LorenzDataset",
        "file_path": cfg.VALID_FILE_PATH,
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.output_keys,
        "block_size": cfg.VALID_BLOCK_SIZE,
        "stride": 1024,
        "embedding_model": embedding_model,
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": False,
        "shuffle": False,
    },
    "batch_size": cfg.EVAL.batch_size,
    "num_workers": 4,
}

mse_validator = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg,
    ppsci.loss.MSELoss(),
    metric={"MSE": ppsci.metric.MSE()},
    name="MSE_Validator",
)
validator = {mse_validator.name: mse_validator}

3.4.5 可视化器构建

本案例中可以通过构建可视化器在模型评估时将评估结果可视化出来，由于 Transformer 模型的输出数据是预测的编码空间的数据无法直接进行可视化，因此需要额外将输出数据使用 Embedding 网络的 Decoder 模块变换到物理状态空间。

在本文中首先定义了对 Transformer 模型输出数据变换到物理状态空间的代码：

examples/lorenz/train_transformer.py
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def build_embedding_model(embedding_model_path: str) -> ppsci.arch.LorenzEmbedding:
    input_keys = ("states",)
    output_keys = ("pred_states", "recover_states")
    regularization_key = "k_matrix"
    model = ppsci.arch.LorenzEmbedding(input_keys, output_keys + (regularization_key,))
    save_load.load_pretrain(model, embedding_model_path)
    return model


class OutputTransform(object):
    def __init__(self, model: base.Arch):
        self.model = model
        self.model.eval()

    def __call__(self, x: Dict[str, paddle.Tensor]):
        pred_embeds = x["pred_embeds"]
        pred_states = self.model.decoder(pred_embeds)

        return pred_states
examples/lorenz/train_transformer.py
64
65
embedding_model = build_embedding_model(cfg.EMBEDDING_MODEL_PATH)
output_transform = OutputTransform(embedding_model)

可以看到，程序首先载入了训练好的 Embedding 模型，然后在 OutputTransform 的 __call__ 函数内实现了编码向量到物理状态空间的变换。

在定义好了以上代码之后，就可以实现可视化器代码的构建了：

examples/lorenz/train_transformer.py
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
states = mse_validator.data_loader.dataset.data
embedding_data = mse_validator.data_loader.dataset.embedding_data
vis_data = {
    "embeds": embedding_data[: cfg.VIS_DATA_NUMS, :-1, :],
    "states": states[: cfg.VIS_DATA_NUMS, 1:, :],
}

visualizer = {
    "visualize_states": ppsci.visualize.VisualizerScatter3D(
        vis_data,
        {
            "pred_states": lambda d: output_transform(d),
            "states": lambda d: d["states"],
        },
        num_timestamps=1,
        prefix="result_states",
    )
}

首先使用上文中的 mse_validator 中的数据集进行可视化，另外还引入了 vis_data_nums 变量用于控制需要可视化样本的数量。最后通过 VisualizerScatter3D 构建可视化器。

3.4.6 模型训练、评估与可视化

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

examples/lorenz/train_transformer.py
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
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,
    eval_freq=cfg.TRAIN.eval_freq,
    validator=validator,
    visualizer=visualizer,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()
# visualize prediction after finished training
solver.visualize()

4. 完整代码

lorenz/train_enn.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
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Two-stage training
# 1. Train a embedding model by running train_enn.py.
# 2. Load pretrained embedding model and freeze it, then train a transformer model by running train_transformer.py.

# This file is for step1: training a embedding model.
# This file is based on PaddleScience/ppsci API.
from os import path as osp

import hydra
import numpy as np
import paddle
from omegaconf import DictConfig

import ppsci
from ppsci.utils import logger


def get_mean_std(data: np.ndarray):
    mean = np.asarray(
        [np.mean(data[:, :, 0]), np.mean(data[:, :, 1]), np.mean(data[:, :, 2])]
    ).reshape(1, 3)
    std = np.asarray(
        [np.std(data[:, :, 0]), np.std(data[:, :, 1]), np.std(data[:, :, 2])]
    ).reshape(1, 3)
    return mean, std


def train(cfg: DictConfig):
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, f"{cfg.mode}.log"), "info")

    weights = (1.0 * (cfg.TRAIN_BLOCK_SIZE - 1), 1.0e4 * cfg.TRAIN_BLOCK_SIZE)
    regularization_key = "k_matrix"
    # manually build constraint(s)
    train_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.TRAIN_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.TRAIN_BLOCK_SIZE,
            "stride": 16,
            "weight_dict": {
                key: value for key, value in zip(cfg.MODEL.output_keys, weights)
            },
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": True,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "num_workers": 4,
    }

    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg,
        ppsci.loss.MSELossWithL2Decay(
            regularization_dict={
                regularization_key: 1.0e-1 * (cfg.TRAIN_BLOCK_SIZE - 1)
            }
        ),
        {
            key: lambda out, k=key: out[k]
            for key in cfg.MODEL.output_keys + (regularization_key,)
        },
        name="Sup",
    )
    constraint = {sup_constraint.name: sup_constraint}

    # set iters_per_epoch by dataloader length
    ITERS_PER_EPOCH = len(sup_constraint.data_loader)

    # manually init model
    data_mean, data_std = get_mean_std(sup_constraint.data_loader.dataset.data)
    model = ppsci.arch.LorenzEmbedding(
        cfg.MODEL.input_keys,
        cfg.MODEL.output_keys + (regularization_key,),
        data_mean,
        data_std,
    )

    # init optimizer and lr scheduler
    clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=0.1)
    lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(
        iters_per_epoch=ITERS_PER_EPOCH,
        decay_steps=ITERS_PER_EPOCH,
        **cfg.TRAIN.lr_scheduler,
    )()
    optimizer = ppsci.optimizer.Adam(
        lr_scheduler, grad_clip=clip, **cfg.TRAIN.optimizer
    )(model)

    # manually build validator
    weights = (1.0 * (cfg.VALID_BLOCK_SIZE - 1), 1.0e4 * cfg.VALID_BLOCK_SIZE)
    eval_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.VALID_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.VALID_BLOCK_SIZE,
            "stride": 32,
            "weight_dict": {
                key: value for key, value in zip(cfg.MODEL.output_keys, weights)
            },
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": False,
        },
        "batch_size": cfg.EVAL.batch_size,
        "num_workers": 4,
    }

    mse_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        ppsci.loss.MSELoss(),
        metric={"MSE": ppsci.metric.MSE()},
        name="MSE_Validator",
    )
    validator = {mse_validator.name: mse_validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        cfg.output_dir,
        optimizer,
        lr_scheduler,
        cfg.TRAIN.epochs,
        ITERS_PER_EPOCH,
        eval_during_train=True,
        validator=validator,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()


def evaluate(cfg: DictConfig):
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, f"{cfg.mode}.log"), "info")

    weights = (1.0 * (cfg.TRAIN_BLOCK_SIZE - 1), 1.0e4 * cfg.TRAIN_BLOCK_SIZE)
    regularization_key = "k_matrix"
    # manually build constraint(s)
    train_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.TRAIN_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.TRAIN_BLOCK_SIZE,
            "stride": 16,
            "weight_dict": {
                key: value for key, value in zip(cfg.MODEL.output_keys, weights)
            },
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": True,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "num_workers": 4,
    }

    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg,
        ppsci.loss.MSELossWithL2Decay(
            regularization_dict={
                regularization_key: 1.0e-1 * (cfg.TRAIN_BLOCK_SIZE - 1)
            }
        ),
        {
            key: lambda out, k=key: out[k]
            for key in cfg.MODEL.output_keys + (regularization_key,)
        },
        name="Sup",
    )

    # manually init model
    data_mean, data_std = get_mean_std(sup_constraint.data_loader.dataset.data)
    model = ppsci.arch.LorenzEmbedding(
        cfg.MODEL.input_keys,
        cfg.MODEL.output_keys + (regularization_key,),
        data_mean,
        data_std,
    )

    # manually build validator
    weights = (1.0 * (cfg.VALID_BLOCK_SIZE - 1), 1.0e4 * cfg.VALID_BLOCK_SIZE)
    eval_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.VALID_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.VALID_BLOCK_SIZE,
            "stride": 32,
            "weight_dict": {
                key: value for key, value in zip(cfg.MODEL.output_keys, weights)
            },
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": False,
        },
        "batch_size": cfg.EVAL.batch_size,
        "num_workers": 4,
    }

    mse_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        ppsci.loss.MSELoss(),
        metric={"MSE": ppsci.metric.MSE()},
        name="MSE_Validator",
    )
    validator = {mse_validator.name: mse_validator}

    solver = ppsci.solver.Solver(
        model,
        output_dir=cfg.output_dir,
        validator=validator,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )
    solver.eval()


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


if __name__ == "__main__":
    main()
lorenz/train_transformer.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
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Two-stage training
# 1. Train a embedding model by running train_enn.py.
# 2. Load pretrained embedding model and freeze it, then train a transformer model by running train_transformer.py.

# This file is for step2: training a transformer model, based on frozen pretrained embedding model.
# This file is based on PaddleScience/ppsci API.
from os import path as osp
from typing import Dict

import hydra
import paddle
from omegaconf import DictConfig

import ppsci
from ppsci.arch import base
from ppsci.utils import logger
from ppsci.utils import save_load


def build_embedding_model(embedding_model_path: str) -> ppsci.arch.LorenzEmbedding:
    input_keys = ("states",)
    output_keys = ("pred_states", "recover_states")
    regularization_key = "k_matrix"
    model = ppsci.arch.LorenzEmbedding(input_keys, output_keys + (regularization_key,))
    save_load.load_pretrain(model, embedding_model_path)
    return model


class OutputTransform(object):
    def __init__(self, model: base.Arch):
        self.model = model
        self.model.eval()

    def __call__(self, x: Dict[str, paddle.Tensor]):
        pred_embeds = x["pred_embeds"]
        pred_states = self.model.decoder(pred_embeds)

        return pred_states


def train(cfg: DictConfig):
    # train time-series: 2048    time-steps: 256    block-size: 64  stride: 64
    # valid time-series: 64      time-steps: 1024   block-size: 256 stride: 1024
    # test  time-series: 256     time-steps: 1024
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, f"{cfg.mode}.log"), "info")

    embedding_model = build_embedding_model(cfg.EMBEDDING_MODEL_PATH)
    output_transform = OutputTransform(embedding_model)

    # manually build constraint(s)
    train_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "file_path": cfg.TRAIN_FILE_PATH,
            "block_size": cfg.TRAIN_BLOCK_SIZE,
            "stride": 64,
            "embedding_model": embedding_model,
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": True,
        },
        "batch_size": cfg.TRAIN.batch_size,
        "num_workers": 4,
    }

    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg,
        ppsci.loss.MSELoss(),
        name="Sup",
    )
    constraint = {sup_constraint.name: sup_constraint}

    # set iters_per_epoch by dataloader length
    ITERS_PER_EPOCH = len(constraint["Sup"].data_loader)

    # manually init model
    model = ppsci.arch.PhysformerGPT2(**cfg.MODEL)

    # init optimizer and lr scheduler
    clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=0.1)
    lr_scheduler = ppsci.optimizer.lr_scheduler.CosineWarmRestarts(
        iters_per_epoch=ITERS_PER_EPOCH, **cfg.TRAIN.lr_scheduler
    )()
    optimizer = ppsci.optimizer.Adam(
        lr_scheduler, grad_clip=clip, **cfg.TRAIN.optimizer
    )(model)

    # manually build validator
    eval_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.VALID_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.VALID_BLOCK_SIZE,
            "stride": 1024,
            "embedding_model": embedding_model,
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": False,
        },
        "batch_size": cfg.EVAL.batch_size,
        "num_workers": 4,
    }

    mse_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        ppsci.loss.MSELoss(),
        metric={"MSE": ppsci.metric.MSE()},
        name="MSE_Validator",
    )
    validator = {mse_validator.name: mse_validator}

    # set visualizer(optional)
    states = mse_validator.data_loader.dataset.data
    embedding_data = mse_validator.data_loader.dataset.embedding_data
    vis_data = {
        "embeds": embedding_data[: cfg.VIS_DATA_NUMS, :-1, :],
        "states": states[: cfg.VIS_DATA_NUMS, 1:, :],
    }

    visualizer = {
        "visualize_states": ppsci.visualize.VisualizerScatter3D(
            vis_data,
            {
                "pred_states": lambda d: output_transform(d),
                "states": lambda d: d["states"],
            },
            num_timestamps=1,
            prefix="result_states",
        )
    }

    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,
        eval_freq=cfg.TRAIN.eval_freq,
        validator=validator,
        visualizer=visualizer,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()
    # visualize prediction after finished training
    solver.visualize()


def evaluate(cfg: DictConfig):
    # directly evaluate pretrained model(optional)
    logger.init_logger("ppsci", osp.join(cfg.output_dir, f"{cfg.mode}.log"), "info")

    embedding_model = build_embedding_model(cfg.EMBEDDING_MODEL_PATH)
    output_transform = OutputTransform(embedding_model)

    # manually init model
    model = ppsci.arch.PhysformerGPT2(**cfg.MODEL)

    # manually build validator
    eval_dataloader_cfg = {
        "dataset": {
            "name": "LorenzDataset",
            "file_path": cfg.VALID_FILE_PATH,
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.output_keys,
            "block_size": cfg.VALID_BLOCK_SIZE,
            "stride": 1024,
            "embedding_model": embedding_model,
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": False,
            "shuffle": False,
        },
        "batch_size": cfg.EVAL.batch_size,
        "num_workers": 4,
    }

    mse_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        ppsci.loss.MSELoss(),
        metric={"MSE": ppsci.metric.MSE()},
        name="MSE_Validator",
    )
    validator = {mse_validator.name: mse_validator}

    # set visualizer(optional)
    states = mse_validator.data_loader.dataset.data
    embedding_data = mse_validator.data_loader.dataset.embedding_data
    vis_datas = {
        "embeds": embedding_data[: cfg.VIS_DATA_NUMS, :-1, :],
        "states": states[: cfg.VIS_DATA_NUMS, 1:, :],
    }

    visualizer = {
        "visulzie_states": ppsci.visualize.VisualizerScatter3D(
            vis_datas,
            {
                "pred_states": lambda d: output_transform(d),
                "states": lambda d: d["states"],
            },
            num_timestamps=1,
            prefix="result_states",
        )
    }

    solver = ppsci.solver.Solver(
        model,
        output_dir=cfg.output_dir,
        validator=validator,
        visualizer=visualizer,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )
    solver.eval()
    # visualize prediction for pretrained model(optional)
    solver.visualize()


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


if __name__ == "__main__":
    main()

5. 结果展示

下图中展示了两个不同初始条件下的模型预测结果和传统数值微分的预测结果。

result_states0

模型预测结果（"pred_states"）与传统数值微分结果（"states"）

result_states1

模型预测结果（"pred_states"）与传统数值微分结果（"states"）
最后更新: November 24, 2023
创建日期: November 6, 2023