TrafNet: Tensor Representation-based Dynamic Graph Neural Network for Traffic Flow Prediction

Abstract: Accurately predicting traffic flow is paramount in addressing congestion issues within urban traffic management. However, traditional deep learning methods face limitations in handling the complex dynamic relationships among multi-source data, coupled with large model parameter counts, high computational complexity, and constraints imposed by purely data-driven approaches. To address these challenges, this study introduces the Tensor Representation-based Auxiliary Information Fusion Network (TrafNet). TrafNet integrates various types of traffic data to construct dynamic graph tensors, utilizing dynamic graph convolution to uncover local dynamic correlations across multi-source data. Furthermore, it enhances global dynamic relationship modeling through shared periodic embeddings, enabling the model to more accurately capture temporal dependencies between traffic data. Additionally, TrafNet employs tensor representation learning to decompose dynamic graph tensors into a multiplicative form of multiple small factors, thereby reducing model parameter counts. Lastly, the introduction of Laplacian graph embeddings as initial parameter values for constructing dynamic graph tensor factors enhances model stability and convergence speed. Experimental results demonstrate that TrafNet performs well on three publicly available datasets, achieving higher prediction accuracy and stability compared to traditional methods. Keywords: Traffic forecasting; Tensor representation; Dynamic graph; Multi-source data fusion; Laplacian Eigenmaps

---

TrafNet implements the model described in Tensor representation-based dynamic graph neural network for traffic flow prediction using auxiliary information (Information Fusion).

⚠️ NOTE: This repository includes integrated code for several auxiliary improvements and experimental modules developed in our group. Only the TrafNet core is required to reproduce the main results; the other extra modules are included for research convenience and are disabled by default.

---

Model diagram

Place the paper figure (architecture) in model_diagram/trafnet_arch.png and include it in your repo. Example markdown to display the image:

Quick start

1. Setup environment

git clone <this-repo-url>
cd <this-repo-dir>
python -m venv venv
source venv/bin/activate            # Linux / macOS
# venv\Scripts\activate             # Windows (PowerShell)
pip install -r requirements.txt

2. Prepare datasets

This repository does not include the original PeMS data. Please prepare datasets under data/ following the names used in experiments:

• PEMS04 (also referenced as PEMSD4) — processed dataset widely used in the literature (we used the version distributed with LightCTS implementation).
• PEMS08 (also referenced as PEMSD8) — processed dataset widely used in the literature (we used the version distributed with LightCTS implementation).
• PEMS03 (PEMSD3) — reprocessed by our group from Caltrans PeMS original raw files; this version contains both flow and speed features (we constructed speed where not provided). Place each dataset under data/ with folder names PEMS03, PEMS04, PEMS08 or as .npz files matching loader expectations.

3. Run training / evaluation

A simple training command (example):

python src/train.py --dataset PEMS08 --batch_size 64 --epochs 200

Other examples:

# Train on PEMS04
python src/train.py --dataset PEMS04 --batch_size 64 --epochs 200

# Train on PEMS03
python src/train.py --dataset PEMS03 --batch_size 64 --epochs 200

Command-line arguments (common options):

• --dataset: PEMS03 / PEMS04 / PEMS08
• --batch_size: batch size (default: 64)
• --epochs: number of training epochs (default: 200)

You can inspect train.py for full available flags.

Acknowledgments

• Many thanks to the open-source implementations that helped shape our codebase and experimental protocols, including LightCTS, DMSTGCN, and Xcompression. See their respective repositories for ideas on data handling, baseline setup, and TT-based compression techniques.
• This work was supported by the members and staff of SmartLab307 (Laboratory/Group name), whose discussions and computational resources were invaluable.

Citation

If you use this code in your work, please cite our paper:

Jianli Zhao, Yiran Hua, Huan Huo, Qiuxia Sun, Qing Li, Hailong Zhang, Tensor representation-based dynamic graph neural network for traffic flow prediction using auxiliary information, Information Fusion, Volume 127, Part B, 2026, 103794, ISSN 1566-2535, https://doi.org/10.1016/j.inffus.2025.103794.

Contact

For questions, issues, or requests, please open an issue on this repository or contact:

• Primary implementer: Yiran Hua — yrhua@sdust.edu.cn
• Corresponding author: Qiuxia Sun — qiuxiasun@sdust.edu.cn