Skip to content

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
Pretrained Model Metric
PEMSD4_pretrained_model.pdparams MAE: 21.48; RMSE: 34.06
PEMSD8_pretrained_model.pdparams MAE: 15.57; RMSE: 24.52

1. Background Introduction

Traffic prediction aims to predict future traffic time series conditions (such as traffic flow or traffic speed) by analyzing historical observation data (such as sensor records on traffic networks). As an important component of the Intelligent Transportation System (ITS), the traffic prediction task is the core foundation for realizing smart cities, including active dynamic traffic control and intelligent route guidance, which helps reduce road safety hazards and improve the operational efficiency of urban transportation systems.

TGCN, a Temporal Graph Convolutional Network for traffic flow prediction. Specifically, by modeling the traffic network as graph structure data, the Graph Convolutional Network (GCN) module is used to extract spatial features; by modeling the traffic signal as time series information, the Temporal Convolutional Network (TCN) module is used to capture temporal features. TGCN finally completes the traffic flow prediction task by iteratively executing two modules.

2. Model Principle

This chapter briefly introduces the model principle of TGCN.

2.1 Graph Convolutional Network Module

This module uses a two-layer message passing network to extract spatial features and update node features:

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 Temporal Convolutional Network Module

This module uses a three-layer one-dimensional convolutional network to extract temporal features and update node features:

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 Model Structure

The TGCN model first uses a feature embedding layer to encode the input signal (i.e., traffic flow data of traffic nodes in the past period):

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

Then the model alternately stacks the aforementioned TCN module and GCN module to update node features:

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)

Finally, the model concatenates the initial node features with the inputs of the two GCN modules, and uses a two-layer MLP to obtain the target output (i.e., traffic flow prediction of traffic nodes in the future period):

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. Model Training

3.1 Dataset Introduction

The case uses preprocessed PEMSD4 and PEMSD8 datasets. PEMSD4 is traffic data from the San Francisco Bay Area, selecting traffic data recorded by 307 sensors on 29 roads from January to February 2018. PEMSD8 is traffic data collected by 170 detectors on 8 roads in San Bernardino from July to August 2016.

Both datasets are saved as N x T x 1 matrices, recording traffic data of corresponding traffic nodes and times, where N is the number of traffic nodes and T is the length of the time series. The two datasets are divided into training set, validation set, and test set according to 7:2:1 respectively. The mean and standard deviation of traffic data are pre-calculated in the case for subsequent normalization operations.

3.2 Model Training

3.2.1 Model Construction

This case is implemented based on the TGCN model, expressed in PaddleScience code as follows:

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 Constraint Construction

This case solves the problem based on data-driven methods, so it is necessary to use SupervisedConstraint built in PaddleScience to construct supervised constraints. Before defining constraints, you need to first specify various parameters used for data loading in constraints.

Training set data loading code is as follows:

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,
}

The code for defining supervised constraints is as follows:

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}

The first parameter of SupervisedConstraint is the data loading method, here train_dataloader_cfg defined above is used;

The second parameter is the definition of loss function, here the custom loss function L1_loss is used;

The third parameter is the name of the constraint condition, which is convenient for subsequent indexing. Here it is named train.

3.2.3 Validator Construction

During the training process of this case, the training status of the current model will be evaluated using the validation set at certain training round intervals, and SupervisedValidator is needed to construct the validator.

Validation set data loading code is as follows:

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,
}

The code for defining supervised validator is as follows:

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}

The SupervisedValidator validator is similar to SupervisedConstraint constraint, the difference is that the validator needs to set evaluation metric metric, here the evaluation metrics used are MAE and RMSE.

3.2.4 Learning Rate and Optimizer Construction

The learning rate size used in this case is set to 1e-2. The optimizer uses Adam, expressed in PaddleScience code as follows:

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

3.2.5 Model Training

After completing the above settings, you only need to pass the instantiated objects to ppsci.solver.Solver, and then start training.

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 Model Export

By setting the eval_during_train parameter in ppsci.solver.Solver, the model parameters with the best effect on the validation set can be automatically saved.

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

3.3 Evaluation Model

3.3.1 Validator Construction

Test set data loading code is as follows:

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,
}

The code for defining supervised validator is as follows:

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}

Similar to SupervisedValidator for validation set, the evaluation metrics used here are MAE and RMSE.

3.3.2 Load Model and Evaluate

Set the loading path of pre-trained model parameters and load the model.

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,
)

Instantiate ppsci.solver.Solver, and then start evaluation.

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. Complete Code

Dataset interface:

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

Model structure:

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}

Model training:

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()

Configuration file:

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. Result Display

The table below shows the evaluation results of TGCN on PEMSD4 and PEMSD8 datasets.

Dataset MAE RMSE
PEMSD4 21.48 34.06
PEMSD8 15.57 24.52