D2STGNN: Decoupled Dynamic Spatial-Temporal Graph Neural Network

Run the data processing script to generate training data for a specific dataset. Replace $DATASET_NAME with one of the supported datasets: METR-LA, PEMS-BAY, PEMS04, or PEMS08. The processed data will be stored in datasets/$DATASET_NAME.

python datasets/raw_data/$DATASET_NAME/generate_training_data.py

To prepare the datasets, download the raw_data.zip and sensor_graph.zip files from the provided Google Drive or BaiduYun links. Unzip them into the datasets directory at the project root, replacing any existing raw_data and sensor_graph folders. Finally, create the log and output directories for training and checkpoints.

cd /path/to/project
unzip raw_data.zip -d ./datasets/
unzip sensor_graph.zip -d ./datasets/
rm {sensor_graph.zip,raw_data.zip}
mkdir log output

To use a pretrained model for testing:

1. Download the pretrained model files from the provided Google Drive or BaiduYun links.
2. Place the downloaded files into the output folder.
3. Check the configuration files for your specific dataset in configs/$DATASET_NAME and ensure the startup arguments are set to test mode.
4. Run the training command (as described in the training section) to execute the test.

Start the training pipeline using main.py by specifying the target dataset with the --dataset flag. Example for the METR-LA dataset:

python main.py --dataset=METR-LA

Install the necessary Python dependencies using pip from the provided requirements.txt file.

pip install -r requirements.txt

D2STGNN is configured through YAML files located in the configs/ directory. The script parses these files to set up data loading, model architecture, and training strategies.

Key configuration sections include:

• start_up: Contains device, model_name, load_pkl (boolean to use cached dataloaders), and mode (test, resume, or scratch).
• data_args: Contains data_dir and adj_data_path.
• model_args: Contains hyperparameters for the D2STGNN model, such as batch_size.
• optim_args: Contains training strategy parameters like epochs, patience (for early stopping), cl_epochs, and warm_epochs.

The trainer function (imported from models.trainer) returns an engine object used to execute the training loop, validation, and testing.

Key methods of the engine include:

• engine.train(trainx, trainy, batch_num, _max, _min): Performs a single training step. Returns mae, mape, and rmse.
• engine.eval(device, dataloader, model_name, _max, _min): Evaluates the model on the validation set. Returns mvalid_loss, mvalid_mape, and mvalid_rmse.
• engine.test(model, save_path_resume, device, dataloader, scaler, model_name, save=False, _max, _min, loss, dataset_name): Runs the final testing procedure.
• engine.set_resume_lr_and_cl(resume_epoch, batch_num): Adjusts the learning rate and contrastive learning steps when resuming training.

The main.py script serves as the primary entrypoint for training, resuming, or testing the D2STGNN model. It uses a configuration-driven approach where a YAML file is loaded based on the --dataset argument provided via the command line.

Supported modes (defined in the config file under start_up.mode) are:

• scratch: Start training from scratch.
• resume: Resume training from a previously saved checkpoint (save_path_resume).
• test: Load the best performing model (save_path) and run evaluation/testing.

Note: The script expects configuration files to exist in the configs/ directory named after the dataset (e.g., configs/METR-LA.yaml).

python main.py --dataset METR-LA

The D2STGNN class is the core neural network architecture. It is initialized using keyword arguments derived from the model_args section of the configuration file.

When initializing, the following parameters are automatically injected or required:

• device: The torch device (CPU/GPU).
• num_nodes: Number of nodes in the graph.
• adjs: A list of adjacency tensors.
• adjs_ori: The original adjacency tensor.
• dataset: The name of the dataset being used.