Traffic Wheel 交通轮

依赖环境

支持python 3.10以上版本。

使用conda创建基本环境

conda create -n trafficwheel python=3.10

pip下载安装包

pip install pyyaml tqdm statsmodels h5py kagglehub torch torchvision torchaudio torchdiffeq fastdtw notebook

快速开始

参考baseline.ipynb中的命令执行，或者使用下面的命令：（请确保当前目录为TrafficWheel）

python run.py --model {model_name} --dataset {dataset_name}  --mode {train, test}  --device {cuda:0}

• model_name: 目前支持：DSANET、STGCN、DCRNN、 GWN(GraphWaveNet)、STSGCN、AGCRN、STFGNN、STGODE、STGNCDE、DDGCRN、TWDGCN
• dataset_name目前支持：PEMSD3，PEMSD4、PEMSD7、PEMSD8
• mode：train为训练模型，test为测试模型。测试模型需要在pre-train文件中找到模型的pth存档。
• device: 支持'cpu'、'cuda:0'、‘cuda:1’ ... 取决于机器卡数

测试模型

在实验结束后，模型的存档文件会被保存在 experiments/dataset/训练时间 文件夹下，共有4个文件。

• best_model.pth 保存了使验证集效果最好的checkpoint
• best_test_model.pth 保存了使测试集的效果最好的checkpoint
• DATASET.yaml 存放了训练模型时所使用的参数
• run.log 记录了训练日志。

可以创建pre-train/{dataset_name}文件夹，把整个文件夹的内容拷贝到experiments/dataset/训练时间 文件夹下的内容全部拷贝到pre-train/{dataset_name}里面。之后就可以在命令中调用 --model test进行测试。