Traffic flow prediction method based on feature embedding and space-time multi-head self-attention mechanism

Abstract

The invention discloses a traffic flow prediction method based on feature embedding and a space-time multi-head self-attention mechanism, and relates to the technical field of intelligent traffic and flow prediction, and the method comprises the steps: obtaining historical traffic flow sequences of N sensors at the first T moments in a traffic network; taking historical traffic flow sequences of N sensors in a traffic network at the first T moments as input, and generating features including time and space through a feature embedding layer; the generated features including time and space are input into a stacked space-time coding layer, calculation is carried out through a decoupled time multi-head self-attention mechanism and a space multi-head self-attention mechanism, new time features and space features are obtained, the time features and the space features are fused, and space-time features are obtained; inputting the spatial-temporal features into an MLP module, and further extracting high-level features through multi-layer MLP; and outputting the high-level features through a regression layer to obtain a prediction result. According to the method, the problem that the traditional prediction model is difficult to process the time-space complexity at the same time is effectively solved by utilizing a decoupling type time and space multi-head self-attention mechanism, and the accuracy of traffic flow prediction is improved.

Description

Technical Field

[0001] The present invention relates to the field of intelligent transportation and flow prediction technology, in particular to a traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism. Background Art

[0002] With the acceleration of urbanization and the rapid growth of car ownership, urban traffic congestion is becoming increasingly serious. Accurately predicting traffic flow is of great significance for formulating reasonable traffic management strategies and optimizing traffic resource allocation. Traditional traffic flow prediction methods are mostly based on statistical models, such as time series analysis and ARIMA models, but since these methods often assume that there is a simple linear relationship between data, they perform poorly when processing complex nonlinear data.

[0003] Traditional feature representation methods often fail to fully extract complex features in traffic data. For example, traffic data usually contains a large amount of high-dimensional unstructured information, and it is a challenge to effectively convert this data into a form that is easy to process. In addition, traditional feature engineering methods require manual selection and extraction of features, which is not only time-consuming and labor-intensive, but also difficult to capture the deep correlation in the data.

[0004] Traffic flow not only changes over time, but also shows certain correlations in space. For example, the changes in traffic flow on adjacent sections often affect each other. Existing prediction methods often find it difficult to capture both temporal and spatial correlations. In addition, due to the complexity of the traffic network, traditional prediction methods often ignore the interactions between nodes in the traffic network, which further limits the accuracy of the prediction.

[0005] In summary, traffic prediction requires feature extraction in both the temporal and spatial dimensions, but existing methods are unable to effectively extract temporal and spatial features. In addition, traffic features are often complex, and converting complex traffic features into features that are conducive to model understanding is also a difficult problem. Summary of the invention

[0006] In order to overcome the defects in the above-mentioned prior art, the present invention provides a traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism to improve the accuracy of traffic flow prediction.

[0007] To achieve the above object, the present invention adopts the following technical solutions, including:

[0008] The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism includes the following steps:

[0009] Step 1: Obtain the historical traffic flow sequence of N sensors in the traffic network at the previous T moments;

[0010] Step 2: Take the historical traffic flow sequence of N sensors in the traffic network at the previous T moments as input, and generate features including time and space through the feature embedding layer;

[0011] Step 3: Input the generated features including time and space into the stacked spatiotemporal coding layer, calculate them through the decoupled temporal multi-head self-attention mechanism and the spatial multi-head self-attention mechanism, respectively, obtain the temporal features and spatial features, and then fuse the temporal features and spatial features to obtain the spatiotemporal features;

[0012] Step 4: Input the spatiotemporal features obtained in step 3 into the MLP module, and further extract high-level features through multi-layer MLP;

[0013] Step 5: The high-level features extracted in step 4 are output through the regression layer to obtain the prediction result. Preferably, in step 2, the feature vector X output by the feature embedding layer E for:

[0014] X E =E w +E d +E n +E p

[0015] Among them, E w Used to indicate a day of the week; E d Used to indicate a certain period of time in a day; n A sensor used to represent a traffic network; E p It is a time position code used to indicate the relative position of each time point in the time series.

[0016] Preferably, in step 3, the temporal multi-head self-attention mechanism is calculated as follows:

[0017] For the input feature vector X E Perform time linear transformation to generate query, key and value matrices. The formula is:

[0018]

[0019] in, are the time parameter matrices of query, key, and value matrices respectively; Q t , K t , V t They are the query, key, and value matrices after time linear transformation;

[0020] Calculate each query matrix Q t And all bond matrices K t The specific calculation formula of the attention weight between is:

[0021]

[0022] Where d is the dimension of the query, key, and value matrices;

[0023] By multiplying the attention score by the value matrix V t , obtain the temporal feature X output by the temporal self-attention mechanism t :

[0024] X t =SoftMax(A t )V t

[0025] Among them, SoftMax(·) is the activation function.

[0026] Preferably, in step 3, the spatial multi-head self-attention mechanism is calculated as follows:

[0027] For the input feature vector X E Perform spatial linear transformation to generate query, key, and value matrices. The formula is:

[0028]

[0029] in, are the spatial parameter matrices of query, key, and value matrices respectively; Q S , K S , V S They are the query, key, and value matrices after spatial linear transformation;

[0030] Calculate each query matrix Q S And all bond matrices K S The specific calculation formula of the attention weight between is:

[0031]

[0032] Where d is the dimension of the query, key, and value matrices;

[0033] By multiplying the attention score by the value matrix V S , obtain the spatial feature X output by the spatial self-attention mechanism s :

[0034] X s =SoftMax(A S )V S

[0035] Among them, SoftMax(·) is the activation function.

[0036] Preferably, in step 3, the outputs of the temporal self-attention mechanism and the spatial self-attention mechanism obtained by the spatiotemporal coding layer are concatenated, and the spatiotemporal feature X output by the spatiotemporal coding layer is st for:

[0037] X st =X t +X s

[0038] Among them, X t is the output of the temporal self-attention mechanism; X s is the output of the spatial self-attention mechanism.

[0039] Preferably, in step 4, the spatiotemporal features Xst are passed through the MLP module to extract high-level features

[0040]

[0041] Among them, MLP(·) is a multi-layer perceptron.

[0042] Preferably, in step 5, the high-level features After the regression layer is output, the prediction result Y is obtained:

[0043]

[0044] Among them, FC(·) is a fully connected layer.

[0045] The present invention also provides a traffic flow prediction model based on feature embedding and spatiotemporal multi-head self-attention mechanism, which is applicable to the above-mentioned traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism, and the model includes:

[0046] Feature embedding layer, used to convert traffic data into a feature vector X containing time and space E ;

[0047] The stacked spatiotemporal coding layer module, each of which contains a temporal multi-head self-attention mechanism and a spatial multi-head self-attention mechanism; the temporal multi-head self-attention mechanism is used to calculate the feature vector X E The correlation between different time points in the spatial multi-head self-attention mechanism is used to calculate the feature vector X E The correlation between different nodes, i.e. different sensors, is used to obtain spatial features. The stacked spatiotemporal coding layer module fuses the temporal features and spatial features to obtain the spatiotemporal features X st ;

[0048] The MLP module includes multiple layers of MLP, which is used to analyze the spatiotemporal features X st Extracting high-level features

[0049] Regression layer, used to Output the prediction results.

[0050] Preferably, the stacked space-time coding layer module is stacked with 3 layers of space-time coding layers.

[0051] The present invention also provides an electronic device, which includes a processor, a memory, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the above-mentioned traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism is implemented.

[0052] The advantages of the present invention are:

[0053] (1) The present invention can effectively convert high-dimensional, unstructured traffic data into low-dimensional, dense feature vectors through the designed feature embedding layer, so as to better capture the key information in the traffic data. This method not only reduces the complexity of data processing, but also improves the model's ability to capture complex features.

[0054] (2) By using the decoupled temporal multi-head self-attention mechanism and the spatial multi-head self-attention mechanism, the present invention can process information along the temporal feature dimension and the spatial feature dimension respectively, thereby effectively capturing the temporal dependency and spatial correlation of traffic flow. This mechanism enables the model to handle temporal and spatial complexity at the same time, thus improving the accuracy of prediction.

[0055] (3) The present invention further extracts high-level features by fusing the temporal feature representation and the spatial feature representation and feeding them into a multi-layer MLP layer, which can fully utilize the interaction between temporal and spatial features and improve prediction accuracy. This method can not only handle nonlinear relationships, but also capture deep feature associations. BRIEF DESCRIPTION OF THE DRAWINGS

[0056] Figure 1 This is the overall framework diagram of the traffic flow prediction model based on feature embedding and spatiotemporal multi-head self-attention mechanism of the present invention.

[0057] Figure 2 This is a diagram of the prediction effect of the present invention on a real data set. DETAILED DESCRIPTION

[0058] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0059] Depend on Figure 1 As shown, the present invention provides a traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism, which can improve the accuracy of traffic flow prediction. The main steps are as follows:

[0060] Step 1: Obtain an existing traffic data set from a public website. The data set contains the historical traffic flow sequence of N sensors in the traffic network at the previous T moments. The data set used in the present invention is the PEMS08 data set, which contains 170 node sensors, and each sensor collects information once every 5 minutes.

[0061] Step 2: Take the historical traffic flow sequence of N sensors in the traffic network at the previous T moments as input, and generate temporal and spatial embedding representations through the feature embedding layer.

[0062] Specifically, the embedded features include: feature E for capturing the periodic patterns of different days of the week w , E w It is used to indicate the day of the week; it is used to capture the characteristics of periodic patterns in different time periods (times) of the day. d , E d Used to distinguish different time periods (times); used to encode features of each sensor (node) in the traffic network n , E n It is used to distinguish specific sensors, help the model better understand the mutual influence between different nodes, and capture spatial dependencies; it is used to encode the feature E of time position. p , E p Represents the relative position of each time point in the time series, thereby better capturing the time dependency.

[0063] For the original Transformer, the present invention adds a temporal position code E p To enhance the model's ability to understand time series data.

[0064] The temporal and spatial features output by the feature embedding layer are:

[0065] X E =E w +E d +E n +E p

[0066] The feature embedding method is used to achieve a simple representation of complex traffic data, thereby modeling the periodic information of traffic data.

[0067] Step 3: Generate the feature X containing time and space EInput to the stacked spatiotemporal encoding layer, calculate through the decoupled temporal multi-head self-attention mechanism and the spatial multi-head self-attention mechanism, obtain new temporal feature representation and spatial feature representation, further fuse the temporal feature representation and the spatial feature representation to obtain the spatiotemporal feature representation; specifically, step 3 can be divided into the following steps:

[0068] Step 3.1: embed the low-dimensional, dense feature vector X generated by the feature embedding layer E The input is fed into the stacked spatiotemporal coding layer for feature extraction. Each layer of the stacked spatiotemporal coding layer contains a temporal multi-head self-attention mechanism and a spatial multi-head self-attention mechanism to independently extract temporal features and spatial features. The temporal features are processed by the decoupled temporal multi-head self-attention mechanism, which can independently calculate the correlation between different time points; while the spatial features are processed by the decoupled spatial multi-head self-attention mechanism, which can independently calculate the correlation between different nodes;

[0069] In the present invention, three layers of spatiotemporal coding layers are used to model dynamic long-distance spatiotemporal dependencies and realize important extraction of traffic data features.

[0070] Step 3.2: Perform the following operations in the temporal multi-head self-attention mechanism:

[0071] For the input feature vector X E Perform time linear transformation to generate query, key and value matrices. The formula is:

[0072]

[0073] in, are the time parameter matrices of query, key, and value matrices respectively; Q t , K t , V t They are the query, key, and value matrices after time linear transformation;

[0074] Calculate each query matrix Q t And all bond matrices K t The specific calculation formula of the attention weight between is:

[0075]

[0076] Where d is the dimension of the query, key, and value matrices;

[0077] By multiplying the attention score by the value matrix V t To obtain the output X of the temporal self-attention mechanism t :

[0078] Xt =SoftMax(A t )V t

[0079] Among them, SoftMax(·) is the activation function;

[0080] Step 3.3: Perform the following operations in the spatial multi-head self-attention mechanism:

[0081] For the input feature vector X E Perform spatial linear transformation to generate query, key, and value matrices. The formula is:

[0082]

[0083] in, are the spatial parameter matrices of query, key, and value matrices respectively; Q S , K S , V S They are the query, key, and value matrices after spatial linear transformation;

[0084] Calculate each query matrix Q S And all bond matrices K S The specific calculation formula of the attention weight between is:

[0085]

[0086] Where d is the dimension of the query, key, and value matrices;

[0087] By multiplying the attention score by the value matrix V S To obtain the output X of the spatial self-attention mechanism s :

[0088] X s =SoftMax(A S )V S

[0089] Specifically, in each layer of temporal and spatial multi-head self-attention mechanism, 8 attention heads are used respectively. There is a layer of FeedForward and LayerNorm after each attention mechanism, and residual connection is used, which can make the model converge faster during training.

[0090] Step 3.4: Concatenate the outputs of the temporal self-attention mechanism and the spatial self-attention mechanism obtained by the spatiotemporal coding layer. The spatiotemporal feature X output by the spatiotemporal coding layer st The following formula:

[0091] X st =X t +Xs

[0092] Specifically, by modeling the time dimension and space dimension respectively, the time and space features are well extracted. Finally, the time and space features are merged, and the output of the spatiotemporal encoding layer is X st Includes dynamic and remote temporal and spatial representations.

[0093] Step 4: The obtained spatiotemporal features X st It represents the input to the MLP module, and further extracts high-level features through multi-layer MLP;

[0094] Specifically, the spatiotemporal feature X st The MLP module further extracts high-level features. It is expressed as:

[0095]

[0096] Among them, MLP(·) is a multi-layer perceptron.

[0097] The MLP module of the present invention uses 3 layers of MLP, the activation function of each layer of MLP is the relu function, and the feature dimension is transformed through the MLP module to further extract more advanced features.

[0098] Step 5: Extract the high-level features After the regression layer is output, the prediction result Y is obtained. The calculation formula is as follows:

[0099]

[0100] Among them, FC(·) is a fully connected layer.

[0101] The present invention uses a feature embedding layer, a spatiotemporal coding layer and an MLP module to construct a feature embedding and spatiotemporal multi-head self-attention mechanism network model for traffic flow prediction. By performing feature embedding representation on historical data, the original high-dimensional, unstructured features are converted into feature representations that are easy for the model to understand. The decoupled time and space multi-head self-attention mechanism in the spatiotemporal coding layer solves the problem of difficulty in modeling dynamic, long-distance spatiotemporal dependencies, and uses the MLP module to further extract high-level features, thereby improving the accuracy of traffic flow prediction.

[0102] The present invention has made predictions on the real data set PEMS08, and the prediction results are as follows: Figure 2 The split ratio of the training set, validation set, and test set is 6:2:2, and Adam is used as the training optimizer of the model. The model uses the historical data of the previous 60 minutes to predict the traffic flow in the next 60 minutes. Using one day's data, the actual value and the predicted value are compared. Figure 2 It can be seen that the trends of the predicted value and the true value are consistent, and the model fitting effect is very good even in peak hours. The results show that the present invention greatly improves the effect of traffic prediction.

[0103] The above are only preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism, characterized by: The following steps are involved: Step 1: Obtain the historical traffic flow sequence of N sensors in the traffic network at the previous T moments; Step 2: Take the historical traffic flow sequence of N sensors in the traffic network at the previous T moments as input, and generate features including time and space through the feature embedding layer; Step 3: Input the generated features including time and space into the stacked spatiotemporal coding layer, calculate them through the decoupled temporal multi-head self-attention mechanism and the spatial multi-head self-attention mechanism, respectively, obtain the temporal features and spatial features, and then fuse the temporal features and spatial features to obtain the spatiotemporal features; Step 4: Input the spatiotemporal features obtained in step 3 into the MLP module, and further extract high-level features through multi-layer MLP; Step 5: Output the high-level features extracted in step 4 through the regression layer to obtain the prediction results.

2. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 2, the feature vector X output by the feature embedding layer E for: X E =And w +E d +E n +E p Among them, E w Used to indicate a day of the week; d Used to indicate a period of time within a day; n A sensor used to represent a traffic network; E p It is a time position code used to indicate the relative position of each time point in the time series.

3. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 3, the temporal multi-head self-attention mechanism is calculated as follows: For the input feature vector X E Perform time linear transformation to generate query, key and value matrices. The formula is: in, are the time parameter matrices of query, key, and value matrices respectively; Q t , K t , V t They are the query, key, and value matrices after time linear transformation; Calculate each query matrix Q t And all bond matrices K t The specific calculation formula of the attention weight between is: Where d is the dimension of the query, key, and value matrices; By multiplying the attention score by the value matrix V t , obtain the temporal feature X output by the temporal self-attention mechanism t : X t =SoftMax(A t )V t Among them, SoftMax(·) is the activation function.

4. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 3, the spatial multi-head self-attention mechanism is calculated as follows: For the input feature vector X E Perform spatial linear transformation to generate query, key, and value matrices. The formula is: in, are the spatial parameter matrices of query, key, and value matrices respectively; Q S , K S , V S They are the query, key, and value matrices after spatial linear transformation; Calculate each query matrix Q S And all bond matrices K S The specific calculation formula of the attention weight between is: Where d is the dimension of the query, key, and value matrices; By multiplying the attention score by the value matrix V S , obtain the spatial feature X output by the spatial self-attention mechanism s : X s =SoftMax(A S )V S Among them, SoftMax(·) is the activation function.

5. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 3, the outputs of the temporal self-attention mechanism and the spatial self-attention mechanism obtained by the spatiotemporal coding layer are concatenated, and the spatiotemporal feature X output by the spatiotemporal coding layer is st for: X st =X t +X s Among them, X t is the output of the temporal self-attention mechanism; X s is the output of the spatial self-attention mechanism.

6. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 4, the spatiotemporal feature X st After the MLP module, high-level features are extracted Among them, MLP(·) is a multi-layer perceptron.

7. The traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 1 is characterized in that: In step 5, the high-level features After the regression layer is output, the prediction result Y is obtained: Among them, FC(·) is a fully connected layer.

8. Traffic flow prediction model based on feature embedding and spatiotemporal multi-head self-attention mechanism, characterized by: The method for predicting traffic flow based on feature embedding and spatiotemporal multi-head self-attention mechanism according to any one of claims 1 to 7 is applicable, and the model includes: Feature embedding layer, used to convert traffic data into a feature vector X containing time and space E ; The stacked spatiotemporal coding layer module, each of which contains a temporal multi-head self-attention mechanism and a spatial multi-head self-attention mechanism; the temporal multi-head self-attention mechanism is used to calculate the feature vector X E The correlation between different time points in the spatial multi-head self-attention mechanism is used to calculate the feature vector X E The correlation between different nodes, i.e. different sensors, is used to obtain spatial features. The stacked spatiotemporal coding layer module fuses the temporal features and spatial features to obtain the spatiotemporal features X st ; The MLP module includes multiple layers of MLP, which is used to analyze the spatiotemporal features X st Extracting high-level features Regression layer, used to Output the prediction results.

9. The traffic flow prediction model based on feature embedding and spatiotemporal multi-head self-attention mechanism according to claim 8, characterized in that: The stacked spatiotemporal coding layer module has 3 stacked spatiotemporal coding layers.

10. An electronic device, characterized in that: It includes a processor, a memory, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the traffic flow prediction method based on feature embedding and spatiotemporal multi-head self-attention mechanism described in any one of claims 1 to 7 is implemented.

Images