czzhangheng
/
Project-I
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Project-I/README.md

3.6 KiB
Raw Blame History

STDEN项目

时空扩散方程网络（Spatio-Temporal Diffusion Equation Network）项目，用于时空序列预测任务。

项目结构

Project-I/
├── run.py                 # 主运行文件
├── configs/              # 配置文件目录
│   ├── stde_gt.yaml     # STDE_GT模型配置
│   ├── stde_wrs.yaml    # STDE_WRS模型配置
│   └── stde_zgc.yaml    # STDE_ZGC模型配置
├── dataloader/           # 数据加载器模块
│   ├── __init__.py
│   └── stden_dataloader.py
├── trainer/              # 训练器模块
│   ├── __init__.py
│   └── stden_trainer.py
├── model/                # 模型模块
│   ├── __init__.py
│   ├── stden_model.py
│   ├── stden_supervisor.py
│   ├── diffeq_solver.py
│   └── ode_func.py
├── lib/                  # 工具库
│   ├── __init__.py
│   ├── logger.py
│   ├── utils.py
│   └── metrics.py
├── requirements.txt      # 项目依赖
└── README.md            # 项目说明

快速开始

1. 安装依赖

pip install -r requirements.txt

2. 训练模型

# 训练STDE_GT模型
python run.py --model_name stde_gt --mode train

# 训练STDE_WRS模型
python run.py --model_name stde_wrs --mode train

# 训练STDE_ZGC模型
python run.py --model_name stde_zgc --mode train

3. 评估模型

# 评估STDE_GT模型
python run.py --model_name stde_gt --mode eval --save_pred

# 评估STDE_WRS模型
python run.py --model_name stde_wrs --mode eval --save_pred

# 评估STDE_ZGC模型
python run.py --model_name stde_zgc --mode eval --save_pred

配置说明

项目使用YAML格式的配置文件，主要包含三个部分：

数据配置 (data)

  • dataset_dir: 数据集目录路径
  • batch_size: 训练批处理大小
  • val_batch_size: 验证批处理大小
  • graph_pkl_filename: 传感器图邻接矩阵文件

模型配置 (model)

  • seq_len: 输入序列长度
  • horizon: 预测时间步数
  • input_dim: 输入特征维度
  • output_dim: 输出特征维度
  • latent_dim: 潜在空间维度
  • n_traj_samples: 轨迹采样数量
  • ode_method: ODE求解方法
  • rnn_units: RNN隐藏单元数量
  • gcn_step: 图卷积步数

训练配置 (train)

  • base_lr: 基础学习率
  • epochs: 总训练轮数
  • patience: 早停耐心值
  • optimizer: 优化器类型
  • max_grad_norm: 最大梯度范数

主要特性

  1. 模块化设计: 清晰的数据加载器、训练器、模型分离
  2. 配置驱动: 使用YAML配置文件，易于调整参数
  3. 统一接口: 通过run.py统一调用不同模型
  4. 完整日志: 支持文件和控制台日志输出
  5. TensorBoard支持: 训练过程可视化
  6. 检查点管理: 自动保存和加载最佳模型

支持的模型

  • STDE_GT: 用于北京GM传感器图数据
  • STDE_WRS: 用于WRS传感器图数据
  • STDE_ZGC: 用于ZGC传感器图数据

数据格式

项目支持两种数据格式：

  1. BJ格式: 包含flow.npz文件，适用于北京数据集
  2. 标准格式: 包含train.npz、val.npz、test.npz文件

日志和输出

  • 训练日志保存在logs/目录
  • 模型检查点保存在checkpoints/目录
  • TensorBoard日志保存在runs/目录
  • 预测结果保存在检查点目录的results/子目录

注意事项

  1. 确保数据集目录结构正确
  2. 传感器图文件路径配置正确
  3. 根据硬件配置调整批处理大小
  4. 训练过程中会自动创建必要的目录

许可证

本项目采用MIT许可证，详见LICENSE文件。