跳转至

TGCN

wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/tgcn/tgcn_data.zip
unzip tgcn_data.zip
python run.py data_name=PEMSD8
# python run.py data_name=PEMSD4

wget -c https://paddle-org.bj.bcebos.com/paddlescience/datasets/tgcn/tgcn_data.zip
unzip tgcn_data.zip
wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/tgcn/PEMSD8_pretrained_model.pdparams
python run.py data_name=PEMSD8 mode=eval EVAL.pretrained_model_path=PEMSD8_pretrained_model.pdparams
# wget -c https://paddle-org.bj.bcebos.com/paddlescience/models/tgcn/PEMSD4_pretrained_model.pdparams
# python run.py data_name=PEMSD4 mode=eval EVAL.pretrained_model_path=PEMSD4_pretrained_model.pdparams
预训练模型 指标
PEMSD4_pretrained_model.pdparams MAE: 21.48; RMSE: 34.06
PEMSD8_pretrained_model.pdparams MAE: 15.57; RMSE: 24.52

1. 背景简介

交通预测旨在通过分析历史观测数据（例如，交通网络上的传感器记录）来预测未来的交通时间序列状况（例如，交通流量或交通速度）。作为智能交通系统（ITS）的重要组成部分，交通预测任务是实现智慧城市的核心基础，包括主动动态交通控制和智能路线引导，有助于减少道路安全隐患并提高城市交通系统的运营效率。

TGCN，一种用于交通流量预测的时空图卷积网络（Temporal Graph Convolutional Network）。具体而言，通过将交通网络建模为图结构数据，使用图卷积网络（GCN）模块提取空间特征；通过将交通信号建模为时序信息，使用时间卷积网络（TCN）模块捕获时间特征。TGCN通过迭代执行两个模块，最终完成交通流量预测任务。

2. 模型原理

本章节对 TGCN 的模型原理进行简单的介绍。

2.1 图卷积网络模块

该模块使用两层消息传递网络，提取空间特征更新节点特征：

ppsci/arch/tgcn.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
class graph_conv(nn.Layer):
    def __init__(self, in_dim, out_dim, dropout, num_layer=2):
        super(graph_conv, self).__init__()
        self.mlp = nn.Conv2D(
            (num_layer + 1) * in_dim,
            out_dim,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )
        self.dropout = dropout
        self.num_layer = num_layer

    def forward(self, x, adj):
        # B C N T
        out = [x]
        for _ in range(self.num_layer):
            new_x = pp.matmul(adj, x)
            out.append(new_x)
            x = new_x

        h = pp.concat(out, axis=1)
        h = self.mlp(h)
        h = F.dropout(h, self.dropout, training=self.training)
        return h

2.2 时间卷积网络模块

该模块使用三层一维卷积网络，提取时间特征更新节点特征：

ppsci/arch/tgcn.py
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
class tempol_conv(nn.Layer):
    def __init__(self, in_dim, out_dim, hidden, num_layer=3, k_s=3, alpha=0.1):
        super(tempol_conv, self).__init__()
        self.leakyrelu = nn.LeakyReLU(alpha)
        self.tc_convs = nn.LayerList()
        self.num_layer = num_layer
        for i in range(num_layer):
            in_channels = in_dim if i == 0 else hidden
            self.tc_convs.append(
                nn.Conv2D(
                    in_channels=in_channels,
                    out_channels=hidden,
                    kernel_size=(1, k_s),
                    padding=(0, i + 1),
                    dilation=i + 1,
                    weight_attr=KaimingNormal(),
                )
            )

        self.mlp = nn.Conv2D(
            in_channels=in_dim + hidden * num_layer,
            out_channels=out_dim,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )

    def forward(self, x):
        # B C N T
        x_cat = [x]
        for i in range(self.num_layer):
            x = self.leakyrelu(self.tc_convs[i](x))
            x_cat.append(x)
        tc_out = self.mlp(pp.concat(x_cat, axis=1))
        return tc_out

2.3 TGCN模型结构

TGCN 模型首先使用特征嵌入层对输入信号（即交通节点在过去一段时间内的流量数据）进行编码：

ppsci/arch/tgcn.py
140
141
142
143
144
145
self.emb_conv = nn.Conv2D(
    in_channels=in_dim,
    out_channels=emb_dim,
    kernel_size=(1, 1),
    weight_attr=KaimingNormal(),
)
ppsci/arch/tgcn.py
173
174
175
176
# emb block
x = raw[self.input_keys[0]]
x = x.transpose(perm=[0, 3, 2, 1])  # B in_dim N T
emb_x = self.emb_conv(x)  # B emd_dim N T

然后模型交替堆叠前述 TCN 模块与 GCN 模块，更新节点特征：

ppsci/arch/tgcn.py
147
148
149
150
151
152
153
154
155
156
157
self.tc1_conv = tempol_conv(
    emb_dim, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha
)
self.sc1_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer)
self.bn1 = nn.BatchNorm2D(hidden)

self.tc2_conv = tempol_conv(
    hidden, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha
)
self.sc2_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer)
self.bn2 = nn.BatchNorm2D(hidden)
ppsci/arch/tgcn.py
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
# TC1
tc1_out = self.tc1_conv(emb_x)  # B hidden N T

# SC1
sc1_out = self.sc1_conv(tc1_out, self.adj)  # B hidden N T
sc1_out = sc1_out + tc1_out
sc1_out = self.bn1(sc1_out)

# TC2
tc2_out = self.tc2_conv(sc1_out)  # B hidden N T

# SC2
sc2_out = self.sc2_conv(tc2_out, self.adj)  # B hidden N T
sc2_out = sc2_out + tc2_out
sc2_out = self.bn2(sc2_out)

最后模型将初始节点特征与两个 GCN 模块的输入拼接，使用两层 MLP 得到目标输出（即交通节点在未来一段时间内的流量预测）：

ppsci/arch/tgcn.py
159
160
161
162
163
164
165
166
167
168
169
170
self.end_conv_1 = nn.Conv2D(
    in_channels=emb_dim + hidden + hidden,
    out_channels=2 * hidden,
    kernel_size=(1, 1),
    weight_attr=KaimingNormal(),
)
self.end_conv_2 = nn.Conv2D(
    in_channels=2 * hidden,
    out_channels=label_len,
    kernel_size=(1, input_len),
    weight_attr=KaimingNormal(),
)
ppsci/arch/tgcn.py
194
195
196
197
198
# readout block
x_out = F.relu(pp.concat((emb_x, sc1_out, sc2_out), axis=1))
x_out = F.relu(self.end_conv_1(x_out))
# transform
x_out = self.end_conv_2(x_out)  # B T N 1

3. 模型训练

3.1 数据集介绍

案例中使用了预处理的 PEMSD4 和 PEMSD8 数据集。PEMSD4 为旧金山湾区交通数据，选取 29 条道路上 307 个传感器记录的交通数据，时间为 2018 年 1 月至 2 月。PEMSD8 为圣贝纳迪诺 8 条道路上 170 个检测器收集的交通数据，时间为 2016 年 7 月至 8 月。

两个数据集均被保存为 N x T x 1 的矩阵，记录了相应交通节点与时间的流量数据，其中 N 为交通节点数量，T 为时间序列长度。两个数据集分别按照 7:2:1 划分为训练集、验证集，和测试集。案例中预先计算了流量数据的均值与标准差，用于后续的正则化操作。

3.2 模型训练

3.2.1 模型构建

该案例基于 TGCN 模型实现，用 PaddleScience 代码表示如下：

examples/tgcn/run.py
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
# set model
model = TGCN(
    input_keys=cfg.MODEL.input_keys,
    output_keys=cfg.MODEL.label_keys,
    adj=adj,
    in_dim=cfg.input_dim,
    emb_dim=cfg.emb_dim,
    hidden=cfg.hidden,
    gc_layer=cfg.gc_layer,
    tc_layer=cfg.tc_layer,
    k_s=cfg.tc_kernel_size,
    dropout=cfg.dropout,
    alpha=cfg.leakyrelu_alpha,
    input_len=cfg.input_len,
    label_len=cfg.label_len,
)

3.2.2 约束器构建

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

训练集数据加载的代码如下:

examples/tgcn/run.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# set train dataloader config
train_dataloader_cfg = {
    "dataset": {
        "name": "PEMSDataset",
        "file_path": cfg.data_path,
        "split": "train",
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.label_keys,
        "norm_input": cfg.norm_input,
        "norm_label": cfg.norm_label,
        "input_len": cfg.input_len,
        "label_len": cfg.label_len,
    },
    "sampler": {
        "name": "BatchSampler",
        "drop_last": True,
        "shuffle": True,
    },
    "batch_size": cfg.TRAIN.batch_size,
}

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

examples/tgcn/run.py
31
32
33
34
35
# set constraint
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg, ppsci.loss.L1Loss(), name="train"
)
constraint = {sup_constraint.name: sup_constraint}

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

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

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

3.2.3 评估器构建

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

验证集数据加载的代码如下:

examples/tgcn/run.py
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
# set eval dataloader config
eval_dataloader_cfg = {
    "dataset": {
        "name": "PEMSDataset",
        "file_path": cfg.data_path,
        "split": "val",
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.label_keys,
        "norm_input": cfg.norm_input,
        "norm_label": cfg.norm_label,
        "input_len": cfg.input_len,
        "label_len": cfg.label_len,
    },
    "sampler": {
        "name": "BatchSampler",
    },
    "batch_size": cfg.EVAL.batch_size,
}

定义监督评估器的代码如下：

examples/tgcn/run.py
56
57
58
59
60
61
62
63
# set validator
sup_validator = ppsci.validate.SupervisedValidator(
    eval_dataloader_cfg,
    ppsci.loss.L1Loss(),
    metric={"MAE": ppsci.metric.MAE(), "RMSE": ppsci.metric.RMSE()},
    name="val",
)
validator = {sup_validator.name: sup_validator}

SupervisedValidator 评估器与 SupervisedConstraint 约束器比较相似，不同的是评估器需要设置评价指标 metric，在这里使用的评价指标分别是 MAE 和 RMSE。

3.2.4 学习率与优化器构建

本案例中学习率大小设置为 1e-2，优化器使用 Adam，用 PaddleScience 代码表示如下：

examples/tgcn/run.py
83
84
# init optimizer
optimizer = ppsci.optimizer.Adam(learning_rate=cfg.TRAIN.learning_rate)(model)

3.2.5 模型训练

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

examples/tgcn/run.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
# initialize solver
solver = ppsci.solver.Solver(
    model=model,
    constraint=constraint,
    output_dir=cfg.output_dir,
    optimizer=optimizer,
    epochs=cfg.TRAIN.epochs,
    iters_per_epoch=iters_per_epoch,
    log_freq=cfg.log_freq,
    eval_during_train=True,
    validator=validator,
    pretrained_model_path=cfg.TRAIN.pretrained_model_path,
    eval_with_no_grad=True,
)
# train model
solver.train()

3.2.6 模型导出

通过设置 ppsci.solver.Solver 中的 eval_during_train 参数，可以自动保存在验证集上效果最优的模型参数。

examples/tgcn/run.py
97
eval_during_train=True,

3.3 评估模型

3.3.1 评估器构建

测试集数据加载的代码如下:

examples/tgcn/run.py
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
test_dataloader_cfg = {
    "dataset": {
        "name": "PEMSDataset",
        "file_path": cfg.data_path,
        "split": "test",
        "input_keys": cfg.MODEL.input_keys,
        "label_keys": cfg.MODEL.label_keys,
        "norm_input": cfg.norm_input,
        "norm_label": cfg.norm_label,
        "input_len": cfg.input_len,
        "label_len": cfg.label_len,
    },
    "sampler": {
        "name": "BatchSampler",
    },
    "batch_size": cfg.EVAL.batch_size,
}

定义监督评估器的代码如下：

examples/tgcn/run.py
127
128
129
130
131
132
133
sup_validator = ppsci.validate.SupervisedValidator(
    test_dataloader_cfg,
    ppsci.loss.L1Loss(),
    metric={"MAE": ppsci.metric.MAE(), "RMSE": ppsci.metric.RMSE()},
    name="test",
)
validator = {sup_validator.name: sup_validator}

与验证集的 SupervisedValidator 相似，在这里使用的评价指标分别是 MAE 和 RMSE。

3.3.2 加载模型并进行评估

设置预训练模型参数的加载路径并加载模型。

examples/tgcn/run.py
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
model = TGCN(
    input_keys=cfg.MODEL.input_keys,
    output_keys=cfg.MODEL.label_keys,
    adj=adj,
    in_dim=cfg.input_dim,
    emb_dim=cfg.emb_dim,
    hidden=cfg.hidden,
    gc_layer=cfg.gc_layer,
    tc_layer=cfg.tc_layer,
    k_s=cfg.tc_kernel_size,
    dropout=cfg.dropout,
    alpha=cfg.leakyrelu_alpha,
    input_len=cfg.input_len,
    label_len=cfg.label_len,
)

实例化 ppsci.solver.Solver，然后启动评估。

examples/tgcn/run.py
155
156
157
158
159
160
161
162
163
164
solver = ppsci.solver.Solver(
    model=model,
    output_dir=cfg.output_dir,
    log_freq=cfg.log_freq,
    validator=validator,
    pretrained_model_path=cfg.EVAL.pretrained_model_path,
    eval_with_no_grad=True,
)
# evaluate
solver.eval()

4. 完整代码

数据集接口：

ppsci/data/dataset/pems_dataset.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
import os
from typing import Dict
from typing import Optional
from typing import Tuple

import numpy as np
import pandas as pd
from paddle.io import Dataset
from paddle.vision.transforms import Compose


class StandardScaler:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def transform(self, data):
        return (data - self.mean) / self.std

    def inverse_transform(self, data):
        return (data * self.std) + self.mean


def add_window_horizon(data, in_step=12, out_step=12):
    length = len(data)
    end_index = length - out_step - in_step
    X = []
    Y = []
    for i in range(end_index + 1):
        X.append(data[i : i + in_step])
        Y.append(data[i + in_step : i + in_step + out_step])
    return X, Y


def get_edge_index(file_path, bi=True, reduce="mean"):
    TYPE_DICT = {0: np.int64, 1: np.int64, 2: np.float32}
    df = pd.read_csv(
        os.path.join(file_path, "dist.csv"),
        skiprows=1,
        header=None,
        sep=",",
        dtype=TYPE_DICT,
    )

    edge_index = df.loc[:, [0, 1]].values.T
    edge_attr = df.loc[:, 2].values

    if bi:
        re_edge_index = np.concatenate((edge_index[1:, :], edge_index[:1, :]), axis=0)
        edge_index = np.concatenate((edge_index, re_edge_index), axis=-1)
        edge_attr = np.concatenate((edge_attr, edge_attr), axis=0)

    num = np.max(edge_index) + 1
    adj = np.zeros((num, num), dtype=np.float32)

    if reduce == "sum":
        adj[edge_index[0], edge_index[1]] = 1.0
    elif reduce == "mean":
        adj[edge_index[0], edge_index[1]] = 1.0
        adj = adj / adj.sum(axis=-1)
    else:
        raise ValueError

    return edge_index, edge_attr, adj


class PEMSDataset(Dataset):
    """Dataset class for PEMSD4 and PEMSD8 dataset.

    Args:
        file_path (str): Dataset root path.
        split (str): Dataset split label.
        input_keys (Tuple[str, ...]): A tuple of input keys.
        label_keys (Tuple[str, ...]): A tuple of label keys.
        weight_dict (Optional[Dict[str, float]]): Define the weight of each constraint variable. Defaults to None.
        transforms (Optional[Compose]): Compose object contains sample wise transform(s). Defaults to None.
        norm_input (bool): Whether to normalize the input. Defaults to True.
        norm_label (bool): Whether to normalize the output. Defaults to False.
        input_len (int): The input timesteps. Defaults to 12.
        label_len (int): The output timesteps. Defaults to 12.

    Examples:
        >>> import ppsci
        >>> dataset = ppsci.data.dataset.PEMSDataset(
        ...     "./Data/PEMSD4",
        ...     "train",
        ...     ("input",),
        ...     ("label",),
        ... )  # doctest: +SKIP
    """

    def __init__(
        self,
        file_path: str,
        split: str,
        input_keys: Tuple[str, ...],
        label_keys: Tuple[str, ...],
        weight_dict: Optional[Dict[str, float]] = None,
        transforms: Optional[Compose] = None,
        norm_input: bool = True,
        norm_label: bool = False,
        input_len: int = 12,
        label_len: int = 12,
    ):
        super().__init__()

        self.input_keys = input_keys
        self.label_keys = label_keys
        self.weight_dict = weight_dict

        self.transforms = transforms
        self.norm_input = norm_input
        self.norm_label = norm_label

        data = np.load(os.path.join(file_path, f"{split}.npy")).astype(np.float32)

        self.mean = np.load(os.path.join(file_path, "mean.npy")).astype(np.float32)
        self.std = np.load(os.path.join(file_path, "std.npy")).astype(np.float32)
        self.scaler = StandardScaler(self.mean, self.std)

        X, Y = add_window_horizon(data, input_len, label_len)
        if norm_input:
            X = self.scaler.transform(X)
        if norm_label:
            Y = self.scaler.transform(Y)

        self._len = X.shape[0]

        self.input = {input_keys[0]: X}
        self.label = {label_keys[0]: Y}

        if weight_dict is not None:
            self.weight_dict = {key: np.array(1.0) for key in self.label_keys}
            self.weight_dict.update(weight_dict)
        else:
            self.weight = {}

    def __getitem__(self, idx):
        input_item = {key: value[idx] for key, value in self.input.items()}
        label_item = {key: value[idx] for key, value in self.label.items()}
        weight_item = {key: value[idx] for key, value in self.weight.items()}

        if self.transforms is not None:
            input_item, label_item, weight_item = self.transforms(
                input_item, label_item, weight_item
            )

        return (input_item, label_item, weight_item)

    def __len__(self):
        return self._len

模型结构：

ppsci/arch/tgcn.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
from typing import Tuple

import paddle as pp
import paddle.nn.functional as F
from numpy import ndarray
from paddle import nn
from paddle.nn.initializer import KaimingNormal

from ppsci.arch.base import Arch


class graph_conv(nn.Layer):
    def __init__(self, in_dim, out_dim, dropout, num_layer=2):
        super(graph_conv, self).__init__()
        self.mlp = nn.Conv2D(
            (num_layer + 1) * in_dim,
            out_dim,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )
        self.dropout = dropout
        self.num_layer = num_layer

    def forward(self, x, adj):
        # B C N T
        out = [x]
        for _ in range(self.num_layer):
            new_x = pp.matmul(adj, x)
            out.append(new_x)
            x = new_x

        h = pp.concat(out, axis=1)
        h = self.mlp(h)
        h = F.dropout(h, self.dropout, training=self.training)
        return h


class tempol_conv(nn.Layer):
    def __init__(self, in_dim, out_dim, hidden, num_layer=3, k_s=3, alpha=0.1):
        super(tempol_conv, self).__init__()
        self.leakyrelu = nn.LeakyReLU(alpha)
        self.tc_convs = nn.LayerList()
        self.num_layer = num_layer
        for i in range(num_layer):
            in_channels = in_dim if i == 0 else hidden
            self.tc_convs.append(
                nn.Conv2D(
                    in_channels=in_channels,
                    out_channels=hidden,
                    kernel_size=(1, k_s),
                    padding=(0, i + 1),
                    dilation=i + 1,
                    weight_attr=KaimingNormal(),
                )
            )

        self.mlp = nn.Conv2D(
            in_channels=in_dim + hidden * num_layer,
            out_channels=out_dim,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )

    def forward(self, x):
        # B C N T
        x_cat = [x]
        for i in range(self.num_layer):
            x = self.leakyrelu(self.tc_convs[i](x))
            x_cat.append(x)
        tc_out = self.mlp(pp.concat(x_cat, axis=1))
        return tc_out


class TGCN(Arch):
    """
    TGCN is a class that represents an Temporal Graph Convolutional Network model.

    Args:
        input_keys (Tuple[str, ...]): A tuple of input keys.
        output_keys (Tuple[str, ...]): A tuple of output keys.
        adj (ndarray): The adjacency matrix of the graph.
        in_dim (int): The dimension of the input data.
        emb_dim (int): The dimension of the embedded space.
        hidden (int): The dimension of the latent space.
        gc_layer (int): The number of the graph convolutional layer.
        tc_layer (int): The number of the temporal convolutional layer.
        k_s (int): The kernel size of the temporal convolutional layer.
        dropout (float): The dropout rate.
        alpha (float): The negative slope of LeakyReLU.
        input_len (int): The input timesteps.
        label_len (int): The output timesteps.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> model = ppsci.arch.TGCN(
        ...     input_keys=("input",),
        ...     output_keys=("label",),
        ...     adj=numpy.ones((307, 307), dtype=numpy.float32),
        ...     in_dim=1,
        ...     emb_dim=32
        ...     hidden=64,
        ...     gc_layer=2,
        ...     tc_layer=2
        ...     k_s=3,
        ...     dropout=0.25,
        ...     alpha=0.1,
        ...     input_len=12,
        ...     label_len=12,
        ... )
        >>> input_dict = {"input": paddle.rand([64, 12, 307, 1]),}
        >>> label_dict = model(input_dict)
        >>> print(label_dict["label"].shape)
        [64, 12, 307, 1]
    """

    def __init__(
        self,
        input_keys: Tuple[str, ...],
        output_keys: Tuple[str, ...],
        adj: ndarray,
        in_dim: int,
        emb_dim: int,
        hidden: int,
        gc_layer: int,
        tc_layer: int,
        k_s: int,
        dropout: float,
        alpha: float,
        input_len: int,
        label_len: int,
    ):
        super(TGCN, self).__init__()

        self.input_keys = input_keys
        self.output_keys = output_keys

        self.register_buffer("adj", pp.to_tensor(data=adj))

        self.emb_conv = nn.Conv2D(
            in_channels=in_dim,
            out_channels=emb_dim,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )

        self.tc1_conv = tempol_conv(
            emb_dim, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha
        )
        self.sc1_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer)
        self.bn1 = nn.BatchNorm2D(hidden)

        self.tc2_conv = tempol_conv(
            hidden, hidden, hidden, num_layer=tc_layer, k_s=k_s, alpha=alpha
        )
        self.sc2_conv = graph_conv(hidden, hidden, dropout, num_layer=gc_layer)
        self.bn2 = nn.BatchNorm2D(hidden)

        self.end_conv_1 = nn.Conv2D(
            in_channels=emb_dim + hidden + hidden,
            out_channels=2 * hidden,
            kernel_size=(1, 1),
            weight_attr=KaimingNormal(),
        )
        self.end_conv_2 = nn.Conv2D(
            in_channels=2 * hidden,
            out_channels=label_len,
            kernel_size=(1, input_len),
            weight_attr=KaimingNormal(),
        )

    def forward(self, raw):
        # emb block
        x = raw[self.input_keys[0]]
        x = x.transpose(perm=[0, 3, 2, 1])  # B in_dim N T
        emb_x = self.emb_conv(x)  # B emd_dim N T

        # TC1
        tc1_out = self.tc1_conv(emb_x)  # B hidden N T

        # SC1
        sc1_out = self.sc1_conv(tc1_out, self.adj)  # B hidden N T
        sc1_out = sc1_out + tc1_out
        sc1_out = self.bn1(sc1_out)

        # TC2
        tc2_out = self.tc2_conv(sc1_out)  # B hidden N T

        # SC2
        sc2_out = self.sc2_conv(tc2_out, self.adj)  # B hidden N T
        sc2_out = sc2_out + tc2_out
        sc2_out = self.bn2(sc2_out)

        # readout block
        x_out = F.relu(pp.concat((emb_x, sc1_out, sc2_out), axis=1))
        x_out = F.relu(self.end_conv_1(x_out))
        # transform
        x_out = self.end_conv_2(x_out)  # B T N 1

        return {self.output_keys[0]: x_out}

模型训练：

examples/tgcn/run.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
import hydra
from omegaconf import DictConfig

import ppsci
from ppsci.arch.tgcn import TGCN
from ppsci.data.dataset.pems_dataset import get_edge_index


def train(cfg: DictConfig):
    # set train dataloader config
    train_dataloader_cfg = {
        "dataset": {
            "name": "PEMSDataset",
            "file_path": cfg.data_path,
            "split": "train",
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.label_keys,
            "norm_input": cfg.norm_input,
            "norm_label": cfg.norm_label,
            "input_len": cfg.input_len,
            "label_len": cfg.label_len,
        },
        "sampler": {
            "name": "BatchSampler",
            "drop_last": True,
            "shuffle": True,
        },
        "batch_size": cfg.TRAIN.batch_size,
    }

    # set constraint
    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg, ppsci.loss.L1Loss(), name="train"
    )
    constraint = {sup_constraint.name: sup_constraint}

    # set eval dataloader config
    eval_dataloader_cfg = {
        "dataset": {
            "name": "PEMSDataset",
            "file_path": cfg.data_path,
            "split": "val",
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.label_keys,
            "norm_input": cfg.norm_input,
            "norm_label": cfg.norm_label,
            "input_len": cfg.input_len,
            "label_len": cfg.label_len,
        },
        "sampler": {
            "name": "BatchSampler",
        },
        "batch_size": cfg.EVAL.batch_size,
    }

    # set validator
    sup_validator = ppsci.validate.SupervisedValidator(
        eval_dataloader_cfg,
        ppsci.loss.L1Loss(),
        metric={"MAE": ppsci.metric.MAE(), "RMSE": ppsci.metric.RMSE()},
        name="val",
    )
    validator = {sup_validator.name: sup_validator}

    # get adj
    _, _, adj = get_edge_index(cfg.data_path, reduce=cfg.reduce)
    # set model
    model = TGCN(
        input_keys=cfg.MODEL.input_keys,
        output_keys=cfg.MODEL.label_keys,
        adj=adj,
        in_dim=cfg.input_dim,
        emb_dim=cfg.emb_dim,
        hidden=cfg.hidden,
        gc_layer=cfg.gc_layer,
        tc_layer=cfg.tc_layer,
        k_s=cfg.tc_kernel_size,
        dropout=cfg.dropout,
        alpha=cfg.leakyrelu_alpha,
        input_len=cfg.input_len,
        label_len=cfg.label_len,
    )
    # init optimizer
    optimizer = ppsci.optimizer.Adam(learning_rate=cfg.TRAIN.learning_rate)(model)
    # set iters_per_epoch by dataloader length
    iters_per_epoch = len(sup_constraint.data_loader)

    # initialize solver
    solver = ppsci.solver.Solver(
        model=model,
        constraint=constraint,
        output_dir=cfg.output_dir,
        optimizer=optimizer,
        epochs=cfg.TRAIN.epochs,
        iters_per_epoch=iters_per_epoch,
        log_freq=cfg.log_freq,
        eval_during_train=True,
        validator=validator,
        pretrained_model_path=cfg.TRAIN.pretrained_model_path,
        eval_with_no_grad=True,
    )
    # train model
    solver.train()


def eval(cfg: DictConfig):
    # set eval dataloader config
    test_dataloader_cfg = {
        "dataset": {
            "name": "PEMSDataset",
            "file_path": cfg.data_path,
            "split": "test",
            "input_keys": cfg.MODEL.input_keys,
            "label_keys": cfg.MODEL.label_keys,
            "norm_input": cfg.norm_input,
            "norm_label": cfg.norm_label,
            "input_len": cfg.input_len,
            "label_len": cfg.label_len,
        },
        "sampler": {
            "name": "BatchSampler",
        },
        "batch_size": cfg.EVAL.batch_size,
    }

    # set validator
    sup_validator = ppsci.validate.SupervisedValidator(
        test_dataloader_cfg,
        ppsci.loss.L1Loss(),
        metric={"MAE": ppsci.metric.MAE(), "RMSE": ppsci.metric.RMSE()},
        name="test",
    )
    validator = {sup_validator.name: sup_validator}

    # get adj
    _, _, adj = get_edge_index(cfg.data_path, reduce=cfg.reduce)
    # set model
    model = TGCN(
        input_keys=cfg.MODEL.input_keys,
        output_keys=cfg.MODEL.label_keys,
        adj=adj,
        in_dim=cfg.input_dim,
        emb_dim=cfg.emb_dim,
        hidden=cfg.hidden,
        gc_layer=cfg.gc_layer,
        tc_layer=cfg.tc_layer,
        k_s=cfg.tc_kernel_size,
        dropout=cfg.dropout,
        alpha=cfg.leakyrelu_alpha,
        input_len=cfg.input_len,
        label_len=cfg.label_len,
    )

    # initialize solver
    solver = ppsci.solver.Solver(
        model=model,
        output_dir=cfg.output_dir,
        log_freq=cfg.log_freq,
        validator=validator,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
        eval_with_no_grad=True,
    )
    # evaluate
    solver.eval()


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


if __name__ == "__main__":
    main()

配置文件：

examples/tgcn/conf/run.yaml
 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
defaults:
  - ppsci_default
  - TRAIN: train_default
  - TRAIN/ema: ema_default
  - TRAIN/swa: swa_default
  - EVAL: eval_default
  - INFER: infer_default
  - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default
  - _self_


hydra:
  run:
    # dynamic output directory according to running time and override name
    dir: outputs_tgcn/${now:%Y-%m-%d}/${now:%H-%M-%S}
  job:
    name: ${mode} # name of logfile
    chdir: false # keep current working directory unchanged
  callbacks:
    init_callback:
      _target_: ppsci.utils.callbacks.InitCallback
  sweep:
    # output directory for multirun
    dir: ${hydra.run.dir}
    subdir: ./

# general settings
device: gpu
mode: train
output_dir: ${hydra:run.dir}
log_freq: 100

# task settings
data_name: PEMSD8
data_path: ./Data/${data_name}
input_len: 12
label_len: 12
norm_input: True
norm_label: False
reduce: mean

# model settings
MODEL:
  input_keys: ["input"]
  label_keys: ["label"]

seed: 3407
batch_size: 64

input_dim: 1
output_dim: 1
emb_dim: 32
hidden: 64
gc_layer: 2
tc_layer: 2
tc_kernel_size: 3
dropout: 0.25
leakyrelu_alpha: 0.1

# training settings
TRAIN:
  epochs: 200
  learning_rate: 0.01
  pretrained_model_path: null
  batch_size: ${batch_size}

# evaluation settings
EVAL:
  pretrained_model_path: null
  batch_size: ${batch_size}

5. 结果展示

下表展示了 TGCN 在 PEMSD4 和 PEMSD8 两个数据集上的评估结果。

数据集 MAE RMSE
PEMSD4 21.48 34.06
PEMSD8 15.57 24.52