This is a pytorch implementation of the SARN paper:
@inproceedings{ChangTC023,
author = {Yanchuan Chang and Egemen Tanin and Xin Cao and Jianzhong Qi},
title = {Spatial Structure-Aware Road Network Embedding via Graph Contrastive Learning},
booktitle = {Proceedings 26th International Conference on Extending Database Technology, {EDBT}},
pages = {144--156},
year = {2023},
}
pip install -r requirements.txt
unzip -j SARN_dataset.zip -d data
First pre-train a SARN (cf. Section Self-supervised Training), then it can be used in downstream tasks, where the parameters can be fine-tuned or frozen (cf. Section Downstream Task Prediction).
Pre-train a SARN model and evaluate its performance on road property prediction task with the frozen embeddings. The trained SARN and the corresponding learned embeddings are persisted to disk (in ./exp/snapshots/
) for other downstream tasks. Logs are ouputted to the terminal and dumped to the log files in ./exp/log
.
python train.py --task_encoder_model SARN --dataset SF
We focus on three downstream tasks on road networks, including road property prediction, trajectory similarity prediction and shortest-path distance prediction - classify, trajsimi and spd for short, respectively.
Fine-tune the pre-trained SARN model and train a classify task model. (Prerequisites: a pre-trained SARN model, in other words, run the last command first).
python train.py --task_encoder_model SARN_ft --dataset SF --task_name classify --task_pretrained_model
Fine-tune the pre-trained SARN model and train a trajsimi task model.
python train.py --task_encoder_model SARN_ft --dataset SF --task_name trajsimi --task_pretrained_model
Fine-tune the pre-trained SARN model and train a spd task model.
python train.py --task_encoder_model SARN_ft --dataset SF --task_name spd --task_pretrained_model
Train a classify task model with the frozen embeddings.
python train.py --task_encoder_model SARN --dataset SF --task_name classify --task_pretrained_model
Train a trajsimi task model with the frozen embeddings.
python train.py --task_encoder_model SARN --dataset SF --task_name trajsimi --task_pretrained_model
Train a spd task model with the frozen embeddings.
python train.py --task_encoder_model SARN --dataset SF --task_name spd --task_pretrained_model
To use your own datasets, you may need to follow the steps below:
./data/OSM_SanFrancisco_downtown_raw
)../utils/osm2roadnetwork.py
and ./data/OSM_SanFrancisco_downtown_raw
)../utils/traj_preprocess_sf.py
). ./task/classifier.py::Classifier.classifier_datasets
, ./task/traj_simi_v2.py::TrajSimi.load_trajsimi_dataset
and ./task/spd.py::SPD.get_spd_dict
).Email changyanchuan@gmail.com if you have any inquiry.