1.5 KiB
1.5 KiB
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进行测试。