Home

Awesome

Traffic Prediction with a graph-based U-Net architecture

This repo contains code to train and evaluate models to predict traffic volue and speed on the data of the Traffic4Cast challenge [1, 2]. [PDF Version of the Paper] [PDF Version of the Slides]

demo_six_cities_2019-04-24_frame156_traffic_sum-1 Sum over 24 consecutive frames of the traffic movies for six different cities on 24.04.2019 starting at 13:00. The color values are scaled logarithmically for better visibility. PDF Version

Abstract

The IARAI Traffic4Cast challenge aims at predicting future traffic in the scale of whole cities. The competitions held in past years have shown that U-Net models perform well on this task. In this work, we point out the advantages of applying graph neural networks instead of visual convolutions and propose a U-Net-like model with graph layers. We further specialize existing graph operations to be sensitive to geographical topology and generalize pooling and upsampling operations to be applicable to graphs. We finally show that our model generalizes well to unseen cities.

Core Concept

As deep convolutional models encode the full input image, later layers will learn patterns specific to the cities in the training set. Intuitively, this limits spatial generalization capabilities. To remove this limitation, we replace the convolutions in UNet with graph layers. The given graph can be described as an incomplete pixel grid. Thus, when comparing two nodes A and B, we can assign a geographic relationship between A and B. For example, node A may be positioned to the south of node B. As edges only ever connect adjacent nodes, there are eight possible values that this relationshoip can have: north (N), north-east (NE), east (E), south-east (SE), south (S), south-west (SW), west (W), north-west (NW). There are two options of incorporating this information in graph models. First, by using edge feature-vectors that specify this geographic relation. Second, by adapting the neighborhood accumulation function in graph layers. We choose the second option and use different learnable transformations depending on the relative position of the neighboring node. Instead of differentiating between all eight, we only differentiate the quadrant NE, SE, SW, NW, as this reflects the way the traffic information is encoded and also reduces the number of transformations. This is tecnically realized by subdividing the graph into four subgraphs that only contain the respective edge direction.

road_subgraphs-1 The four directed subgraphs that we extracted from the road graph. The graphs exclusively contain edges in the respective geographical direction.

Subgraph creation is done in data/data_utils. The graph operation is defined in layers/hybrid_gnn.py. As you can see, we iterate over the subgraphs (s. line 85) and apply separate dense layers in every iteration. This ultimately leads to output node features that are sensitive to the geographical neighborhood topology.

What's contained in this repository

Setup and Usage of the Code

Get the data and setup the environment

  1. Download the dataset from the competition (unfortunately the link cannot be published here).
  2. Extract the data anywhere you like
  3. Install the dependencies: pip install -r requirements.txt (we use Python 3.8.10)

Train a model

  --model 
      : Select the model to train (s. config.py for a list of possible models to train)
  --activation, --depth, --units, --use_bias
      : Set model hyperparameters with these arguments
  --layer_type
      : (Only graph-models) Used to specify a layer type for graph-models
  --batch, --epochs,
      : Configure the training loop
  --acc_gradients_steps 
      : Set the number of gradient accumulation steps before updating the model
  --data_type, --data_dir 
      : Configure the dataset

Examples

Train a vanilla UNet model with eight down- and upsampling blocks, respectively. Note that the data in this case is located in data/raw:

python train.py --model UNet \
             --depth 8 --units 8 \
             --batch 2 --epochs 15 --acc_gradients_steps 8 \
             --data_type 'image' --data_dir data/raw

Train a ReLU-activated Graph-UNet model from the provided checkpoint:

python train.py --model GraphUNet \
             --batch 1 --epochs 15 --acc_gradients_steps 16 \
             --data_type 'graph' --data_dir data/raw \
             --checkpoint ckpts/GraphUNet/GraphUNet_03-10-2021__16-04-37

When first training a model that uses graph data, all dataset files are preprocessed to speed up the data pipeline

Create a submission

Example

python submission.py --model GraphUNet \
             --batch 1 --layer_type 'geo_quadrant_gcn' \
             --data_type 'graph' --data_dir data/raw \
             --checkpoint ckpts/GraphUNet/GraphUNet_03-10-2021__16-04-37/

Reproduce the Results

Here, n_days is set to 3 to make the script run quicker. In our evaluations, we set n_days to 30 to reduce the variance in the results. The figures will be saved in results/<model_id>.

python results.py --model GraphUNet \
             --batch 1 --layer_type 'geo_quadrant_gcn' \
             --data_type 'graph' --data_dir data/raw \
             --checkpoint ckpts/GraphUNet/GraphUNet_03-10-2021__16-04-37/ \
             --n_days 3

Citation

@INPROCEEDINGS{hermes_graphunet_2022,
  author={Hermes, Luca and Hammer, Barbara and Melnik, Andrew and Velioglu, Riza and Vieth, Markus and Schilling, Malte},
  booktitle={2022 International Joint Conference on Neural Networks (IJCNN)},
  title={A Graph-based U-Net Model for Predicting Traffic in unseen Cities},
  year={2022},
  pages={1-8},
  doi={10.1109/IJCNN55064.2022.9892453}
}

References