impl mtgnn

config/MTGNN/AirQuality.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: AirQuality
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 6
+  lag: 24
+  normalizer: std
+  num_nodes: 35
+  steps_per_day: 24
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 35  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 6 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 6 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 6
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/BJTaxi-Inflow.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: BJTaxi-InFlow
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 1024
+  steps_per_day: 48
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 1024  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/BJTaxi-Outflow.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: BJTaxi-OutFlow
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 1024
+  steps_per_day: 48
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 1024  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/METR-LA.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: METR-LA
+  device: cuda:1
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 207
+  steps_per_day: 288
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 207 # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/NYCBike-Inflow.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: NYCBike-InFlow
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 128
+  steps_per_day: 48
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 128  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/NYCBike-Outflow.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: NYCBike-OutFlow
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 128
+  steps_per_day: 48
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 128  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/PEMS-BAY.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: PEMS-BAY
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 1
+  lag: 24
+  normalizer: std
+  num_nodes: 325
+  steps_per_day: 288
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 325  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/MTGNN/SolarEnergy.yaml

@@ -0,0 +1,64 @@
+basic:
+  dataset: SolarEnergy
+  device: cuda:0
+  mode: train
+  model: MTGNN
+  seed: 2023
+
+data:
+  batch_size: 64
+  column_wise: false
+  days_per_week: 7
+  horizon: 24
+  input_dim: 6
+  lag: 24
+  normalizer: std
+  num_nodes: 137
+  steps_per_day: 24
+  test_ratio: 0.2
+  val_ratio: 0.2
+
+model:
+  gcn_true: True  # 是否使用图卷积网络 (bool)
+  buildA_true: True  # 是否动态构建邻接矩阵 (bool)
+  subgraph_size: 20  # 子图大小 (int)
+  num_nodes: 137  # 节点数量 (int)
+  node_dim: 40  # 节点嵌入维度 (int)
+  dilation_exponential: 1  # 膨胀卷积指数 (int)
+  conv_channels: 32  # 卷积通道数 (int)
+  residual_channels: 32  # 残差通道数 (int)
+  skip_channels: 64  # 跳跃连接通道数 (int)
+  end_channels: 128  # 输出层通道数 (int)
+  seq_len: 24  # 输入序列长度 (int)
+  in_dim: 1 # 输入特征维度 (int)
+  out_len: 24 # 输出序列长度 (int)
+  out_dim: 1 # 输出预测维度 (int)
+  layers: 3  # 模型层数 (int)
+  propalpha: 0.05 # 图传播参数alpha (float)
+  tanhalpha: 3 # tanh激活参数alpha (float)
+  layer_norm_affline: True # 层归一化是否使用affine变换 (bool)
+  gcn_depth: 2  # 图卷积深度 (int)
+  dropout: 0.3  # dropout率 (float)
+  predefined_A: null # 预定义邻接矩阵 (optional, None)
+  static_feat: null  # 静态特征 (optional, None)
+
+train:
+  batch_size: 64
+  debug: false
+  early_stop: true
+  early_stop_patience: 15
+  epochs: 100
+  grad_norm: false
+  log_step: 1000
+  loss_func: mae
+  lr_decay: true
+  lr_decay_rate: 0.3
+  lr_decay_step: 5,20,40,70
+  lr_init: 0.003
+  mae_thresh: None
+  mape_thresh: 0.001
+  max_grad_norm: 5
+  output_dim: 1
+  plot: false
+  real_value: true
+  weight_decay: 0

config/PatchTST/AirQuality.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: AirQuality
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 6
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/BJTaxi-Inflow.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: BJTaxi-InFlow
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/BJTaxi-Outflow.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: BJTaxi-OutFlow
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/METR-LA.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: METR-LA
   device: cuda:1
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/NYCBike-Inflow.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: NYCBike-InFlow
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/NYCBike-Outflow.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: NYCBike-OutFlow
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -23,6 +23,7 @@ model:
   seq_len: 24
   pred_len: 24
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_model: 128
   d_ff: 2048

config/PatchTST/PEMS-BAY.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: PEMS-BAY
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -24,6 +24,7 @@ model:
   pred_len: 24
   d_model: 128
   patch_len: 6
+  enc_in: 1
   stride: 8
   d_ff: 2048
   dropout: 0.1

config/PatchTST/SolarEnergy.yaml

@@ -2,7 +2,7 @@ basic:
   dataset: SolarEnergy
   device: cuda:0
   mode: train
-  model: iTransformer
+  model: PatchTST
   seed: 2023
 
 data:
@@ -24,6 +24,7 @@ model:
   pred_len: 24
   d_model: 128
   patch_len: 6
+  enc_in: 6
   stride: 8
   d_ff: 2048
   dropout: 0.1

model/MTGNN/MTGNN.py

@@ -3,91 +3,109 @@ from model.MTGNN.layer import *
 
 
 class gtnet(nn.Module):
-    def __init__(self, gcn_true, buildA_true, gcn_depth, num_nodes, device, predefined_A=None, static_feat=None, dropout=0.3, subgraph_size=20, node_dim=40, dilation_exponential=1, conv_channels=32, residual_channels=32, skip_channels=64, end_channels=128, seq_length=12, in_dim=2, out_dim=12, layers=3, propalpha=0.05, tanhalpha=3, layer_norm_affline=True):
+    def __init__(self, configs):
         super(gtnet, self).__init__()
-        self.gcn_true = gcn_true
+        self.gcn_true = configs['gcn_true']  # 是否使用图卷积网络
-        self.buildA_true = buildA_true
+        self.buildA_true = configs['buildA_true']  # 是否动态构建邻接矩阵
-        self.num_nodes = num_nodes
+        self.num_nodes = configs['num_nodes']  # 节点数量
-        self.dropout = dropout
+        self.device = configs['device']  # 设备（CPU/GPU）
-        self.predefined_A = predefined_A
+        self.dropout = configs['dropout']  # dropout率
-        self.filter_convs = nn.ModuleList()
+        self.predefined_A = configs.get('predefined_A', None)  # 预定义邻接矩阵
-        self.gate_convs = nn.ModuleList()
+        self.static_feat = configs.get('static_feat', None)  # 静态特征
-        self.residual_convs = nn.ModuleList()
+        self.subgraph_size = configs['subgraph_size']  # 子图大小
-        self.skip_convs = nn.ModuleList()
+        self.node_dim = configs['node_dim']  # 节点嵌入维度
-        self.gconv1 = nn.ModuleList()
+        self.dilation_exponential = configs['dilation_exponential']  # 膨胀卷积指数
-        self.gconv2 = nn.ModuleList()
+        self.conv_channels = configs['conv_channels']  # 卷积通道数
-        self.norm = nn.ModuleList()
+        self.residual_channels = configs['residual_channels']  # 残差通道数
-        self.start_conv = nn.Conv2d(in_channels=in_dim,
+        self.skip_channels = configs['skip_channels']  # 跳跃连接通道数
-                                    out_channels=residual_channels,
+        self.end_channels = configs['end_channels']  # 输出层通道数
+        self.seq_length = configs['seq_len']  # 输入序列长度
+        self.in_dim = configs['in_dim']  # 输入特征维度
+        self.out_len = configs['out_len']  # 输出序列长度
+        self.out_dim = configs['out_dim']  # 输出预测维度
+        self.layers = configs['layers']  # 模型层数
+        self.propalpha = configs['propalpha']  # 图传播参数alpha
+        self.tanhalpha = configs['tanhalpha']  # tanh激活参数alpha
+        self.layer_norm_affline = configs['layer_norm_affline']  # 层归一化是否使用affine变换
+        self.gcn_depth = configs['gcn_depth']  # 图卷积深度
+        self.filter_convs = nn.ModuleList()  # 卷积滤波器列表
+        self.gate_convs = nn.ModuleList()  # 门控卷积列表
+        self.residual_convs = nn.ModuleList()  # 残差卷积列表
+        self.skip_convs = nn.ModuleList()  # 跳跃连接卷积列表
+        self.gconv1 = nn.ModuleList()  # 第一层图卷积列表
+        self.gconv2 = nn.ModuleList()  # 第二层图卷积列表
+        self.norm = nn.ModuleList()  # 归一化层列表
+        self.start_conv = nn.Conv2d(in_channels=self.in_dim,
+                                    out_channels=self.residual_channels,
                                     kernel_size=(1, 1))
-        self.gc = graph_constructor(num_nodes, subgraph_size, node_dim, device, alpha=tanhalpha, static_feat=static_feat)
+        self.gc = graph_constructor(self.num_nodes, self.subgraph_size, self.node_dim, self.device, alpha=self.tanhalpha, static_feat=self.static_feat)
 
-        self.seq_length = seq_length
         kernel_size = 7
-        if dilation_exponential>1:
+        if self.dilation_exponential>1:
-            self.receptive_field = int(1+(kernel_size-1)*(dilation_exponential**layers-1)/(dilation_exponential-1))
+            self.receptive_field = int(1+(kernel_size-1)*(self.dilation_exponential**self.layers-1)/(self.dilation_exponential-1))
         else:
-            self.receptive_field = layers*(kernel_size-1) + 1
+            self.receptive_field = self.layers*(kernel_size-1) + 1
 
         for i in range(1):
-            if dilation_exponential>1:
+            if self.dilation_exponential>1:
-                rf_size_i = int(1 + i*(kernel_size-1)*(dilation_exponential**layers-1)/(dilation_exponential-1))
+                rf_size_i = int(1 + i*(kernel_size-1)*(self.dilation_exponential**self.layers-1)/(self.dilation_exponential-1))
             else:
-                rf_size_i = i*layers*(kernel_size-1)+1
+                rf_size_i = i*self.layers*(kernel_size-1)+1
             new_dilation = 1
-            for j in range(1,layers+1):
+            for j in range(1,self.layers+1):
-                if dilation_exponential > 1:
+                if self.dilation_exponential > 1:
-                    rf_size_j = int(rf_size_i + (kernel_size-1)*(dilation_exponential**j-1)/(dilation_exponential-1))
+                    rf_size_j = int(rf_size_i + (kernel_size-1)*(self.dilation_exponential**j-1)/(self.dilation_exponential-1))
                 else:
                     rf_size_j = rf_size_i+j*(kernel_size-1)
 
-                self.filter_convs.append(dilated_inception(residual_channels, conv_channels, dilation_factor=new_dilation))
+                self.filter_convs.append(dilated_inception(self.residual_channels, self.conv_channels, dilation_factor=new_dilation))
-                self.gate_convs.append(dilated_inception(residual_channels, conv_channels, dilation_factor=new_dilation))
+                self.gate_convs.append(dilated_inception(self.residual_channels, self.conv_channels, dilation_factor=new_dilation))
-                self.residual_convs.append(nn.Conv2d(in_channels=conv_channels,
+                self.residual_convs.append(nn.Conv2d(in_channels=self.conv_channels,
-                                                    out_channels=residual_channels,
+                                                    out_channels=self.residual_channels,
                                                  kernel_size=(1, 1)))
                 if self.seq_length>self.receptive_field:
-                    self.skip_convs.append(nn.Conv2d(in_channels=conv_channels,
+                    self.skip_convs.append(nn.Conv2d(in_channels=self.conv_channels,
-                                                    out_channels=skip_channels,
+                                                    out_channels=self.skip_channels,
                                                     kernel_size=(1, self.seq_length-rf_size_j+1)))
                 else:
-                    self.skip_convs.append(nn.Conv2d(in_channels=conv_channels,
+                    self.skip_convs.append(nn.Conv2d(in_channels=self.conv_channels,
-                                                    out_channels=skip_channels,
+                                                    out_channels=self.skip_channels,
                                                     kernel_size=(1, self.receptive_field-rf_size_j+1)))
 
                 if self.gcn_true:
-                    self.gconv1.append(mixprop(conv_channels, residual_channels, gcn_depth, dropout, propalpha))
+                    self.gconv1.append(mixprop(self.conv_channels, self.residual_channels, self.gcn_depth, self.dropout, self.propalpha))
-                    self.gconv2.append(mixprop(conv_channels, residual_channels, gcn_depth, dropout, propalpha))
+                    self.gconv2.append(mixprop(self.conv_channels, self.residual_channels, self.gcn_depth, self.dropout, self.propalpha))
 
                 if self.seq_length>self.receptive_field:
-                    self.norm.append(LayerNorm((residual_channels, num_nodes, self.seq_length - rf_size_j + 1),elementwise_affine=layer_norm_affline))
+                    self.norm.append(LayerNorm((self.residual_channels, self.num_nodes, self.seq_length - rf_size_j + 1),elementwise_affine=self.layer_norm_affline))
                 else:
-                    self.norm.append(LayerNorm((residual_channels, num_nodes, self.receptive_field - rf_size_j + 1),elementwise_affine=layer_norm_affline))
+                    self.norm.append(LayerNorm((self.residual_channels, self.num_nodes, self.receptive_field - rf_size_j + 1),elementwise_affine=self.layer_norm_affline))
 
-                new_dilation *= dilation_exponential
+                new_dilation *= self.dilation_exponential
 
-        self.layers = layers
+        self.end_conv_1 = nn.Conv2d(in_channels=self.skip_channels,
-        self.end_conv_1 = nn.Conv2d(in_channels=skip_channels,
+                                             out_channels=self.end_channels,
-                                             out_channels=end_channels,
                                              kernel_size=(1,1),
                                              bias=True)
-        self.end_conv_2 = nn.Conv2d(in_channels=end_channels,
+        self.end_conv_2 = nn.Conv2d(in_channels=self.end_channels,
-                                             out_channels=out_dim,
+                                             out_channels=self.out_len * self.out_dim,
                                              kernel_size=(1,1),
                                              bias=True)
         if self.seq_length > self.receptive_field:
-            self.skip0 = nn.Conv2d(in_channels=in_dim, out_channels=skip_channels, kernel_size=(1, self.seq_length), bias=True)
+            self.skip0 = nn.Conv2d(in_channels=self.in_dim, out_channels=self.skip_channels, kernel_size=(1, self.seq_length), bias=True)
-            self.skipE = nn.Conv2d(in_channels=residual_channels, out_channels=skip_channels, kernel_size=(1, self.seq_length-self.receptive_field+1), bias=True)
+            self.skipE = nn.Conv2d(in_channels=self.residual_channels, out_channels=self.skip_channels, kernel_size=(1, self.seq_length-self.receptive_field+1), bias=True)
+
 
         else:
-            self.skip0 = nn.Conv2d(in_channels=in_dim, out_channels=skip_channels, kernel_size=(1, self.receptive_field), bias=True)
+            self.skip0 = nn.Conv2d(in_channels=self.in_dim, out_channels=self.skip_channels, kernel_size=(1, self.receptive_field), bias=True)
-            self.skipE = nn.Conv2d(in_channels=residual_channels, out_channels=skip_channels, kernel_size=(1, 1), bias=True)
+            self.skipE = nn.Conv2d(in_channels=self.residual_channels, out_channels=self.skip_channels, kernel_size=(1, 1), bias=True)
 
-
+        self.idx = torch.arange(self.num_nodes).to(self.device)
-        self.idx = torch.arange(self.num_nodes).to(device)
 
 
     def forward(self, input, idx=None):
+        input = input[..., :-2] # 去掉周期嵌入
+        input = input.transpose(1, 3)
         seq_len = input.size(3)
         assert seq_len==self.seq_length, 'input sequence length not equal to preset sequence length'
 
@@ -130,5 +148,8 @@ class gtnet(nn.Module):
         skip = self.skipE(x) + skip
         x = F.relu(skip)
         x = F.relu(self.end_conv_1(x))
-        x = self.end_conv_2(x)
+        x = self.end_conv_2(x) # [b, t*c, n, 1]
+        # [b, t*c, n, 1] -> [b,t,c,n] -> [b, t, n, c]
+        x = x.reshape(x.size(0), self.out_len, self.out_dim, self.num_nodes)
+        x = x.permute(0, 1, 3, 2)
         return x

model/PatchTST/PatchTST.py

@@ -62,14 +62,14 @@ class Model(nn.Module):
                     activation=configs['activation']
                 ) for l in range(configs['e_layers'])
             ],
-            norm_layer=nn.Sequential(Transpose(1,2), nn.BatchNorm1d(configs.d_model), Transpose(1,2))
+            norm_layer=nn.Sequential(Transpose(1,2), nn.BatchNorm1d(configs['d_model']), Transpose(1,2))
         )
 
         # Prediction Head
-        self.head_nf = configs.d_model * \
+        self.head_nf = configs['d_model'] * \
-                       int((configs.seq_len - self.patch_len) / self.stride + 2)
+                       int((configs['seq_len'] - self.patch_len) / self.stride + 2)
-        self.head = FlattenHead(configs.enc_in, self.head_nf, configs.pred_len,
+        self.head = FlattenHead(configs['enc_in'], self.head_nf, configs['pred_len'],
-                                head_dropout=configs.dropout)
+                                head_dropout=configs['dropout'])
 
     def forecast(self, x_enc):
         # Normalization from Non-stationary Transformer

model/PatchTST/layers/Embed.py

@@ -1,5 +1,26 @@
 import torch
 import torch.nn as nn
+import math
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEmbedding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float()
+                    * -(math.log(10000.0) / d_model)).exp()
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe[:, :x.size(1)]
 
 class PatchEmbedding(nn.Module):
     def __init__(self, d_model, patch_len, stride, padding, dropout):

model/model_selector.py

@@ -30,6 +30,7 @@ from model.ASTRA.astrav3 import ASTRA as ASTRAv3
 from model.iTransformer.iTransformer import iTransformer
 from model.HI.HI import HI
 from model.PatchTST.PatchTST import Model as PatchTST
+from model.MTGNN.MTGNN import gtnet as MTGNN
 
 
 
@@ -99,3 +100,5 @@ def model_selector(config):
             return HI(model_config)
         case "PatchTST":
             return PatchTST(model_config)
+        case "MTGNN":
+            return MTGNN(model_config)

train.py

@@ -1,10 +1,27 @@
 import yaml
 import torch
+import os
 
 import utils.initializer as init
 from dataloader.loader_selector import get_dataloader
 from trainer.trainer_selector import select_trainer
 
+def read_config(config_path):
+    with open(config_path, "r") as file:
+        config = yaml.safe_load(file)
+    
+    # 全局配置
+    device = "cuda:0" # 指定设备
+    seed = 2023 # 随机种子
+    epochs = 100
+
+    # 拷贝项
+    config["basic"]["device"] = device
+    config["model"]["device"] = device
+    config["basic"]["seed"] = seed
+    config["train"]["epochs"] = epochs
+    return config
+
 def run(config):
     init.init_seed(config["basic"]["seed"])
     model = init.init_model(config)
@@ -45,22 +62,26 @@ def run(config):
 
 if __name__ == "__main__":
     # 指定模型
-    model_list = ["PatchTST"]
+    model_list = ["MTGNN"]
     # 指定数据集
-    dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]
+    dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]    
     # dataset_list = ["AirQuality"]
-    device = "cuda:0" # 指定设备
+
-    seed = 2023 # 随机种子
+    # 我的调试开关，不做测试就填 str(False)
-    epochs = 1
+    os.environ["TRY"] = str(False)
+    
     for model in model_list:
         for dataset in dataset_list:
             config_path = f"./config/{model}/{dataset}.yaml"
-            with open(config_path, "r") as file:
+            # 可去这个函数里面调整统一的config项，⚠️注意调设备，epochs
-                config = yaml.safe_load(file)
+            config = read_config(config_path)
-            config["basic"]["device"] = device
+            print(f"\nRunning {model} on {dataset}")
-            config["basic"]["seed"] = seed
+            # print(f"config: {config}")
-            config["train"]["epochs"] = epochs
+            if os.environ.get("TRY") == "True":
-            print(f"\nRunning {model} on {dataset} with seed {seed} on {device}")
+                try:
-            print(f"config: {config}")
+                    run(config)
-            run(config)
+                except Exception as e:
+                    pass
+            else:
+                run(config)
 

trainer/Trainer.py

@@ -71,6 +71,10 @@ class Trainer:
                 label = target[..., : self.args["output_dim"]]
                 # 计算loss和反归一化loss
                 output = self.model(data)
+                # 我的调试开关
+                if os.environ.get("TRY") == "True":
+                    print(f"[{'✅' if output.shape == label.shape else '❌'}]: output: {output.shape}, label: {label.shape}")
+                    assert False
                 loss = self.loss(output, label)
                 d_output = self.scaler.inverse_transform(output)
                 d_label = self.scaler.inverse_transform(label)