A Traffic Flow Prediction Framework Based on Clustering and Heterogeneous Graph Neural Networks

Traffic flow forecasting is crucial for traffic management, but the complex spatio-temporal correlation and heterogeneity among traffic nodes make this problem challenging. While many deep spatio-temporal models have been proposed and applied to traffic flow prediction, they mostly focus on capturing the spatio-temporal correlation among traffic nodes, ignoring the influence of the functional characteristics of the area to which the nodes belong. Therefore, there is a need to propose a method to help models capture such influence. This paper presents a novel framework that enhances existing deep spatio-temporal models by combining clustering with heterogeneous graph neural networks. Our framework's clustering module measures the similarity between nodes in the traffic pattern using the Dynamic Time Warping and the Wasserstein distance and then applies spectral clustering to divide the nodes into different clusters based on traffic pattern. Our graph transformer module can adaptively construct a new graph for nodes in the same cluster, and the spatio-temporal feature learning module captures the spatio-temporal correlation among nodes based on the new graph. Extensive experiments on two real datasets demonstrate that our proposed framework can effectively improve the performance of some representative deep spatio-temporal models.