Traffic Prediction using Recurrent Temporal GNN
In this tutorial, we will learn how to use a recurrent Temporal Graph Convolutional Network (TGCN) to predict traffic in a spatio-temporal setting. Traffic forecasting is the problem of predicting future traffic trends on a road network given historical traffic data, such as, in our case, traffic speed and time of day.
Import
We start by importing the necessary libraries. We use GNNLux.jl, Lux.jl and MLDatasets.jl, among others.
using Lux, GNNLux
using MLDatasets: METRLA
using Statistics, Plots, Random
using Zygote, Optimisers
const mae = MAELoss()
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
rng = Random.seed!(42); # for reproducibilityDataset: METR-LA
We use the METR-LA dataset from the paper Diffusion Convolutional Recurrent Neural Network: Data-driven Traffic Forecasting, which contains traffic data from loop detectors in the highway of Los Angeles County. The dataset contains traffic speed data from March 1, 2012 to June 30, 2012. The data is collected every 5 minutes, resulting in 12 observations per hour, from 207 sensors. Each sensor is a node in the graph, and the edge weights are the distances between the sensor locations.
dataset_metrla = METRLA(; num_timesteps = 3)dataset METRLA:
graphs => 1-element Vector{MLDatasets.Graph}g = dataset_metrla[1]Graph:
num_nodes => 207
num_edges => 1722
edge_index => ("1722-element Vector{Int64}", "1722-element Vector{Int64}")
node_data => (features = "34269-element Vector{Any}", targets = "34269-element Vector{Any}")
edge_data => 1722-element Vector{Float32}edge_data contains the weights of the edges of the graph and node_data contains a node feature vector and a target vector. The latter vectors contain batches of dimension num_timesteps, which means that they contain vectors with the node features and targets of num_timesteps time steps. Two consecutive batches are shifted by one-time step. The node features are the traffic speed of the sensors and the time of the day, and the targets are the traffic speed of the sensors in the next time step. Let's see some examples:
features = map(x -> permutedims(x,(1,3,2)), g.node_data.features)
size(features[1])(2, 3, 207)The first dimension corresponds to the two features (the first line the speed value and the second line the time of day), the second to the number of time steps num_timesteps and the third to the nodes.
targets = map(x -> permutedims(x,(1,3,2)), g.node_data.targets)
size(targets[1])(1, 3, 207)In the case of the targets the first dimension is 1 because they store just the speed value.
features[1][:,:,1]2×3 Matrix{Float32}:
1.17081 1.11647 1.15888
-0.876741 -0.87663 -0.87652features[2][:,:,1]2×3 Matrix{Float32}:
1.11647 1.15888 -0.876741
-0.87663 -0.87652 -0.87641targets[1][:,:,1]1×3 Matrix{Float32}:
1.11647 1.15888 -0.876741function plot_data(data,sensor)
p = plot(legend=false, xlabel="Time (h)", ylabel="Normalized speed")
plotdata = []
for i in 1:3:length(data)
push!(plotdata,data[i][1,:,sensor])
end
plotdata = reduce(vcat,plotdata)
plot!(p, collect(1:length(data)), plotdata, color = :green, xticks =([i for i in 0:50:250], ["$(i)" for i in 0:4:24]))
return p
end
plot_data(features[1:288],1) # Plot the speed of the first sensor for the first day
Now let's construct the static graph, the train_loader and data_loader.
graph = GNNGraph(g.edge_index; edata = g.edge_data, g.num_nodes);
train_loader = zip(features[1:288], targets[1:288]); # train on 24 hours
test_loader = zip(features[289:577], targets[289:577]); # test on next 24 hoursModel: T-GCN
We use the T-GCN model from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction] (https://arxiv.org/pdf/1811.05320.pdf), which consists of a graph convolutional network (GCN) and a gated recurrent unit (GRU). The GCN is used to capture spatial features from the graph, and the GRU is used to capture temporal features from the feature time series.
model = GNNChain(TGCN(2 => 100; add_self_loops = false), Dense(100, 1))GNNChain(
layers = NamedTuple(
layer_1 = GNNRecurrence(
cell = TGCNCell(
conv_z = GNNChain(
layers = NamedTuple(
layer_1 = GCNConv(2 => 100, relu, add_self_loops=false), # 300 parameters
layer_2 = GCNConv(100 => 100, add_self_loops=false), # 10_100 parameters
),
),
dense_z = Dense(200 => 100, σ), # 20_100 parameters
conv_r = GNNChain(
layers = NamedTuple(
layer_1 = GCNConv(2 => 100, relu, add_self_loops=false), # 300 parameters
layer_2 = GCNConv(100 => 100, add_self_loops=false), # 10_100 parameters
),
),
dense_r = Dense(200 => 100, σ), # 20_100 parameters
conv_h = GNNChain(
layers = NamedTuple(
layer_1 = GCNConv(2 => 100, relu, add_self_loops=false), # 300 parameters
layer_2 = GCNConv(100 => 100, add_self_loops=false), # 10_100 parameters
),
),
dense_h = Dense(200 => 100, tanh), # 20_100 parameters
),
),
layer_2 = Dense(100 => 1), # 101 parameters
),
) # Total: 91_601 parameters,
# plus 0 states.Let's look at the output of the model for the first batch of the training data. In Lux, the parameters and the state of the model are kept separate from the model itself, so we first initialize them.
ps, st = Lux.setup(rng, model);
ŷ, _ = model(graph, features[1], ps, st);
ŷ1×3×207 Array{Float32, 3}:
[:, :, 1] =
-0.00821716 0.0144252 0.0331109
[:, :, 2] =
-0.00965475 0.0114244 0.0313205
[:, :, 3] =
-0.021871 -0.00552126 0.00915654
[:, :, 4] =
-0.0011948 0.0242512 0.0433345
[:, :, 5] =
-0.00177842 0.0234212 0.0419707
[:, :, 6] =
-0.00148941 0.0233941 0.0428406
[:, :, 7] =
0.000171542 0.0263602 0.0475252
[:, :, 8] =
-0.0239704 -0.00853027 0.004878
[:, :, 9] =
-0.0168092 -0.000113703 0.0207041
[:, :, 10] =
-0.00818147 0.0140459 0.0358767
[:, :, 11] =
-0.00618431 0.0190886 0.0378305
[:, :, 12] =
-0.0309838 -0.0105868 0.0095427
[:, :, 13] =
-0.000436462 0.0265467 0.0472731
[:, :, 14] =
-0.00863666 0.0150596 0.0368315
[:, :, 15] =
-0.0117289 0.00825214 0.0348723
[:, :, 16] =
0.000949625 0.027768 0.0486651
[:, :, 17] =
0.00283761 0.0309462 0.0569695
[:, :, 18] =
-0.00238523 0.022369 0.0403578
[:, :, 19] =
-0.0306872 -0.0087631 -0.00276568
[:, :, 20] =
-0.0059651 0.0184265 0.0383082
[:, :, 21] =
-0.00548585 0.0193472 0.0405481
[:, :, 22] =
-0.0250331 -0.00580959 0.00109873
[:, :, 23] =
-0.00356281 0.0230084 0.0466084
[:, :, 24] =
-0.00714046 0.0130345 0.0312584
[:, :, 25] =
-0.0118014 0.0109802 0.0286891
[:, :, 26] =
-0.0151689 0.00218989 0.0175212
[:, :, 27] =
-0.00242122 0.0275407 0.0453983
[:, :, 28] =
-0.0346418 -0.02133 -0.00626384
[:, :, 29] =
-0.0146368 0.00559739 0.0229257
[:, :, 30] =
-0.0164983 0.000262856 0.0112214
[:, :, 31] =
0.00149859 0.0272778 0.0527369
[:, :, 32] =
-0.0110724 0.0105973 0.0257418
[:, :, 33] =
-0.0106025 0.00989559 0.0268476
[:, :, 34] =
0.00138379 0.0284617 0.049704
[:, :, 35] =
-0.011029 0.00959434 0.0297334
[:, :, 36] =
-0.0277521 -0.00468668 0.00168869
[:, :, 37] =
-0.0106502 0.0124526 0.0259133
[:, :, 38] =
-0.0101708 0.0108753 0.0293083
[:, :, 39] =
-0.00371963 0.0198472 0.0389333
[:, :, 40] =
-0.0087364 0.013785 0.0318417
[:, :, 41] =
-0.013489 0.00636882 0.0191159
[:, :, 42] =
-0.0033318 0.0221838 0.0422297
[:, :, 43] =
-0.0100745 0.0113125 0.0298342
[:, :, 44] =
-0.00307915 0.0233894 0.0496784
[:, :, 45] =
-0.0293883 -0.00140207 0.0153636
[:, :, 46] =
-0.0203904 -0.000731792 0.0178434
[:, :, 47] =
-0.0242228 -0.0108217 7.44835f-5
[:, :, 48] =
-0.00794924 0.0155703 0.0320494
[:, :, 49] =
-0.000443846 0.0238108 0.0461237
[:, :, 50] =
-0.00823136 0.0133004 0.0322236
[:, :, 51] =
-0.0182157 -0.000923406 0.00874105
[:, :, 52] =
-0.0197446 0.0024194 0.0129717
[:, :, 53] =
-0.00978112 0.0111993 0.026973
[:, :, 54] =
-0.0142461 0.00751363 0.025807
[:, :, 55] =
-0.0101709 0.0108753 0.0293082
[:, :, 56] =
-0.0297162 -0.0130692 0.00278344
[:, :, 57] =
-0.0109779 0.00445043 0.0219942
[:, :, 58] =
-0.0100355 0.0117707 0.029253
[:, :, 59] =
-0.0107693 0.0114266 0.0306436
[:, :, 60] =
-0.0184939 -0.00159456 0.0146626
[:, :, 61] =
-0.00412727 0.0208961 0.0406313
[:, :, 62] =
-0.0133903 0.00800258 0.0177869
[:, :, 63] =
-0.0160819 0.00204028 0.0180206
[:, :, 64] =
-0.00412727 0.0208961 0.0406313
[:, :, 65] =
-0.00385667 0.0188402 0.0378552
[:, :, 66] =
-0.00477769 0.01704 0.0354537
[:, :, 67] =
-0.02122 -0.00378448 0.00212817
[:, :, 68] =
-0.0057567 0.0203873 0.0390659
[:, :, 69] =
-0.00440173 0.022102 0.0441296
[:, :, 70] =
-0.0055023 0.0198467 0.0387614
[:, :, 71] =
-0.00894158 0.014867 0.0314622
[:, :, 72] =
-0.00465536 0.0215146 0.0430236
[:, :, 73] =
-0.00324241 0.0248326 0.0482382
[:, :, 74] =
-0.00730579 0.0162243 0.0332827
[:, :, 75] =
-0.0121717 0.00827613 0.0243309
[:, :, 76] =
-0.00820319 0.0151796 0.0315771
[:, :, 77] =
-0.00511628 0.0193203 0.0399505
[:, :, 78] =
-0.00664458 0.0168547 0.0434569
[:, :, 79] =
-0.0212049 -0.00525067 0.00802272
[:, :, 80] =
-0.0326336 -0.0174345 -0.00179935
[:, :, 81] =
0.000316545 0.0269748 0.0474231
[:, :, 82] =
-0.00344989 0.022354 0.0433216
[:, :, 83] =
-0.0018939 0.0250479 0.0508254
[:, :, 84] =
-0.0104273 0.0105077 0.0252024
[:, :, 85] =
-0.0111967 0.0086285 0.0364158
[:, :, 86] =
-0.0359868 -0.0197313 -0.00156351
[:, :, 87] =
-0.00905617 0.0120922 0.0267821
[:, :, 88] =
-0.00705249 0.0131835 0.0311062
[:, :, 89] =
-0.0102068 0.0108329 0.039408
[:, :, 90] =
-0.00732871 0.0155721 0.0379833
[:, :, 91] =
-0.00896793 0.0119808 0.0278258
[:, :, 92] =
0.000171535 0.0263602 0.0475251
[:, :, 93] =
-0.0292947 -0.0123728 0.0047844
[:, :, 94] =
0.0018874 0.0290393 0.0512981
[:, :, 95] =
-0.0095189 0.0127057 0.0287307
[:, :, 96] =
-0.00920017 0.0153237 0.0363948
[:, :, 97] =
-0.00675596 0.0171003 0.0365007
[:, :, 98] =
-0.00516706 0.0195319 0.0380939
[:, :, 99] =
-0.00468567 0.0212621 0.0432734
[:, :, 100] =
-0.0109695 0.0100322 0.0249101
[:, :, 101] =
-0.00480064 0.0175132 0.0372615
[:, :, 102] =
-0.00175633 0.021358 0.0421611
[:, :, 103] =
-0.00675705 0.0183506 0.0370809
[:, :, 104] =
-0.00435386 0.0223792 0.0449273
[:, :, 105] =
-0.00948677 0.012299 0.0277941
[:, :, 106] =
-0.0310867 -0.0244714 -0.0213755
[:, :, 107] =
-0.039139 -0.0349075 -0.0324736
[:, :, 108] =
-0.0197428 -0.00222888 0.0120533
[:, :, 109] =
-0.0340344 -0.0180384 -0.00232398
[:, :, 110] =
-0.016809 0.0020473 0.0137824
[:, :, 111] =
-0.0258126 -0.00946338 0.00629207
[:, :, 112] =
-0.0113044 0.00965049 0.0271328
[:, :, 113] =
-0.01398 0.00985853 0.0250579
[:, :, 114] =
-0.00451003 0.0224715 0.0425556
[:, :, 115] =
-0.0107693 0.0114266 0.0306436
[:, :, 116] =
-0.0128593 0.00807549 0.0214224
[:, :, 117] =
-0.0116444 0.00762197 0.0239993
[:, :, 118] =
-0.0108711 0.0117586 0.0249509
[:, :, 119] =
-0.00560274 0.0204472 0.0397052
[:, :, 120] =
-0.0121718 0.00827614 0.0243309
[:, :, 121] =
-0.0151689 0.00218991 0.0175213
[:, :, 122] =
-0.0158763 0.000667531 0.00323578
[:, :, 123] =
-0.00835298 0.0145499 0.0314063
[:, :, 124] =
-0.0062982 0.0156706 0.0363288
[:, :, 125] =
-0.00961787 0.0109065 0.0273044
[:, :, 126] =
-0.00960421 0.014225 0.036486
[:, :, 127] =
-0.0163094 0.00533274 0.00596636
[:, :, 128] =
-0.00718839 0.0162061 0.0339369
[:, :, 129] =
-0.00301997 0.0229231 0.0425203
[:, :, 130] =
-0.00686035 0.0145443 0.0342574
[:, :, 131] =
-0.00408264 0.0218743 0.0489538
[:, :, 132] =
-0.00594971 0.0171859 0.0346717
[:, :, 133] =
-0.0358658 -0.0192914 -0.000742409
[:, :, 134] =
-0.00103747 0.0219577 0.0432877
[:, :, 135] =
-0.0055554 0.016153 0.0345801
[:, :, 136] =
-0.00702389 0.0144618 0.0335251
[:, :, 137] =
0.000548251 0.0267444 0.0477673
[:, :, 138] =
-0.00526601 0.0187382 0.0375164
[:, :, 139] =
0.000486601 0.0265969 0.0475849
[:, :, 140] =
-0.0104272 0.0093519 0.0272851
[:, :, 141] =
-0.00952506 0.0126065 0.0285398
[:, :, 142] =
-0.00322082 0.0238048 0.0447394
[:, :, 143] =
-0.0092582 0.0113804 0.0304347
[:, :, 144] =
-0.0118315 0.00599926 0.0232591
[:, :, 145] =
0.000520378 0.0270373 0.0474962
[:, :, 146] =
-0.0125182 0.00752925 0.0221856
[:, :, 147] =
-0.00328181 0.0224241 0.0491296
[:, :, 148] =
-0.0108386 0.008702 0.0250677
[:, :, 149] =
-0.00445814 0.0179113 0.0386408
[:, :, 150] =
-0.0143861 0.00381661 0.0164905
[:, :, 151] =
-0.0108416 0.00402115 0.0214087
[:, :, 152] =
-0.0102509 0.0109841 0.0389173
[:, :, 153] =
-0.0270753 0.00281014 0.0234369
[:, :, 154] =
-0.0199297 -0.0041375 0.00616935
[:, :, 155] =
-0.00125065 0.0257128 0.0471909
[:, :, 156] =
-0.00341222 0.0222467 0.0423326
[:, :, 157] =
-0.00917533 0.0137064 0.0298123
[:, :, 158] =
-0.00322511 0.0219762 0.0404697
[:, :, 159] =
-0.006485 0.0166181 0.0342301
[:, :, 160] =
-0.00353946 0.0222017 0.0419642
[:, :, 161] =
-0.00112146 0.0259375 0.0471458
[:, :, 162] =
-0.00145669 0.0243457 0.0434363
[:, :, 163] =
-0.00114095 0.025566 0.0460631
[:, :, 164] =
-0.00399512 0.0216691 0.0413134
[:, :, 165] =
-0.00746722 0.0159911 0.0328996
[:, :, 166] =
-0.00272395 0.0207096 0.0416148
[:, :, 167] =
0.00169669 0.0287229 0.0532165
[:, :, 168] =
-0.0178723 0.00496994 0.0235218
[:, :, 169] =
-0.0116743 0.0119938 0.0318312
[:, :, 170] =
-0.011565 0.0115453 0.0313856
[:, :, 171] =
-0.00521046 0.012802 0.0320369
[:, :, 172] =
-0.0187236 0.00331177 0.0242589
[:, :, 173] =
-0.00695409 0.0158522 0.0344385
[:, :, 174] =
-0.0156356 0.00588394 0.0229958
[:, :, 175] =
-0.00952948 0.0107777 0.0265258
[:, :, 176] =
-0.00528941 0.0203201 0.0417827
[:, :, 177] =
-0.0120879 0.00868366 0.0226738
[:, :, 178] =
-0.00867432 0.0133468 0.0330589
[:, :, 179] =
-0.0271077 0.00287797 0.023712
[:, :, 180] =
-0.0292645 -0.0097343 0.00924958
[:, :, 181] =
-0.00465781 0.0169174 0.0371693
[:, :, 182] =
-0.00607961 0.0166836 0.0353746
[:, :, 183] =
-0.0146344 0.00502735 0.0249425
[:, :, 184] =
-0.0116239 0.00889217 0.0274019
[:, :, 185] =
-0.0197955 -0.00350938 0.00836946
[:, :, 186] =
-0.0142319 0.00364592 0.0281281
[:, :, 187] =
-0.010957 0.00966091 0.0342825
[:, :, 188] =
-0.00322242 0.0222684 0.0412074
[:, :, 189] =
0.00044544 0.0276311 0.0486719
[:, :, 190] =
-0.0188853 0.00264707 0.00703365
[:, :, 191] =
-0.00330989 0.024461 0.0483255
[:, :, 192] =
-0.000522427 0.0256213 0.0452741
[:, :, 193] =
-0.00106541 0.0245034 0.0437235
[:, :, 194] =
-0.00190246 0.0238661 0.0429132
[:, :, 195] =
-0.0144073 0.00319022 0.0169229
[:, :, 196] =
-0.000914615 0.0247917 0.0438946
[:, :, 197] =
0.00169674 0.0287229 0.0532165
[:, :, 198] =
-0.000986643 0.0247489 0.0485976
[:, :, 199] =
-0.00350053 0.0219981 0.0410607
[:, :, 200] =
-0.00952503 0.0126065 0.0285399
[:, :, 201] =
-0.0310937 -0.0111156 -0.00414067
[:, :, 202] =
-0.0221025 0.00454363 0.0246241
[:, :, 203] =
-0.0162078 0.00394783 0.0199544
[:, :, 204] =
-0.0201875 -0.00508371 0.00845758
[:, :, 205] =
-0.00533505 0.0184017 0.037045
[:, :, 206] =
-0.0069662 0.0168798 0.0363927
[:, :, 207] =
-0.00615102 0.0178156 0.0365203The output of the model is a tensor of size (1, 3, 207), which corresponds to the dimension of the feature (in this case speed), the number of time steps, and the number of nodes in the graph, respectively. The model outputs the predicted traffic speed for each sensor at each time step.
Training
We train the model for 100 epochs, using the Adam optimizer with a learning rate of 0.001. We use the mean absolute error (MAE) as the loss function.
function loss_function(model, ps, st, (graph, x, y))
ŷ, st = model(graph, x, ps, st)
return mae(ŷ, y), st, (;)
end
function train(graph, train_loader, model, ps, st)
train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001))
for epoch in 1:100
for (x, y) in train_loader
_, _, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss_function, (graph, x, y), train_state)
end
if epoch % 10 == 0
ps, st = train_state.parameters, train_state.states
loss = mean([mae(model(graph, x, ps, st)[1], y) for (x, y) in train_loader])
@show epoch, loss
end
end
return train_state.parameters, train_state.states
end
ps, st = train(graph, train_loader, model, ps, st);(epoch, loss) = (10, 0.26200253f0)
(epoch, loss) = (20, 0.2570168f0)
(epoch, loss) = (30, 0.253641f0)
(epoch, loss) = (40, 0.2518906f0)
(epoch, loss) = (50, 0.25033492f0)
(epoch, loss) = (60, 0.2518252f0)
(epoch, loss) = (70, 0.25138408f0)
(epoch, loss) = (80, 0.24981792f0)
(epoch, loss) = (90, 0.24679154f0)
(epoch, loss) = (100, 0.2435512f0)
function plot_predicted_data(graph, features, targets, sensor)
p = plot(xlabel="Time (h)", ylabel="Normalized speed")
prediction = []
ground_truth = []
for i in 1:3:length(features)
push!(ground_truth,targets[i][1,:,sensor])
push!(prediction, model(graph, features[i], ps, st)[1][1,:,sensor])
end
prediction = reduce(vcat,prediction)
ground_truth = reduce(vcat, ground_truth)
plot!(p, collect(1:length(prediction)), prediction, color = :red, label= "Prediction")
plot!(p, collect(1:length(ground_truth)), ground_truth, color = :blue, label = "Ground Truth", xticks = ([i for i in 0:50:250], ["$(i)" for i in 0:4:20]))
return p
end
plot_predicted_data(graph,features[289:577],targets[289:577], 1)
accuracy(ŷ, y) = 1 - Statistics.norm(y-ŷ)/Statistics.norm(y)accuracy (generic function with 1 method)Test accuracy:
mean([accuracy(model(graph, x, ps, st)[1], y) for (x, y) in test_loader])0.6099206f0The accuracy is not very good but can be improved by training using more data. We used a small subset of the dataset for this tutorial because of the computational cost of training the model. From the plot of the predictions, we can see that the model is able to capture the general trend of the traffic speed, but it is not able to capture the peaks of the traffic.
Conclusion
In this tutorial, we learned how to use a recurrent temporal graph convolutional network to predict traffic in a spatio-temporal setting. We used the TGCN model, which consists of a graph convolutional network (GCN) and a gated recurrent unit (GRU). We then trained the model for 100 epochs on a small subset of the METR-LA dataset. The accuracy of the model is not very good, but it can be improved by training on more data.
This page was generated using Literate.jl.