framework

data/data_selector.py

@@ -0,0 +1,13 @@
+import numpy as np
+import os
+
+def load_dataset(config):
+    dataset_name = config['basic']['dataset']
+    node_num = config['data']['num_nodes']
+    input_dim = config['data']['input_dim']
+    data = None
+    match dataset_name:
+        case 'EcoSolar':
+            data_path = os.path.join('./data/EcoSolar.npy')
+            data = np.load(data_path)[:, :node_num, :input_dim]
+    return data

data/dataloader.py

@@ -0,0 +1,170 @@
+from utils.normalizer import normalize_dataset
+
+import numpy as np
+import gc
+import torch
+
+def get_dataloader(config, data):
+    config = config['data']
+    L, N, F = data.shape  # 数据形状
+
+    # Step 1: data -> x,y
+    x = add_window_x(data, config['lag'], config['horizon'])
+    y = add_window_y(data, config['lag'], config['horizon'])
+
+    del data
+    gc.collect()
+
+    # Step 2: time_in_day, day_in_week -> day, week
+    time_in_day = [i % config['steps_per_day'] / config['steps_per_day'] for i in range(L)]
+    time_in_day = np.tile(np.array(time_in_day), [1, N, 1]).transpose((2, 1, 0))
+    day_in_week = [(i // config['steps_per_day']) % config['days_per_week'] for i in range(L)]
+    day_in_week = np.tile(np.array(day_in_week), [1, N, 1]).transpose((2, 1, 0))
+
+    x_day = add_window_x(time_in_day, config['lag'], config['horizon'])
+    x_week = add_window_x(day_in_week, config['lag'], config['horizon'])
+
+    # Step 3 day, week, x, y --> x, y
+    x = np.concatenate([x, x_day, x_week], axis=-1)
+
+    del x_day, x_week
+    gc.collect()
+
+    # Step 4 x,y --> x_train, x_val, x_test, y_train, y_val, y_test
+    if config['test_ratio'] > 1:
+        x_train, x_val, x_test = split_data_by_days(x, config['val_ratio'], config['test_ratio'])
+    else:
+        x_train, x_val, x_test = split_data_by_ratio(x, config['val_ratio'], config['test_ratio'])
+
+    del x
+    gc.collect()
+
+    # Multi-channel normalization - each channel normalized independently
+    input_channels = x_train[..., :config['input_dim']]
+    num_channels = input_channels.shape[-1]
+    
+    # Initialize scalers for each channel
+    scalers = []
+    
+    # Normalize each channel independently
+    for channel_idx in range(num_channels):
+        channel_data = input_channels[..., channel_idx:channel_idx+1]
+        scaler = normalize_dataset(channel_data, config['normalizer'], config['column_wise'])
+        scalers.append(scaler)
+        
+        # Apply transformation to each channel
+        x_train[..., channel_idx:channel_idx+1] = scaler.transform(channel_data)
+        x_val[..., channel_idx:channel_idx+1] = scaler.transform(x_val[..., channel_idx:channel_idx+1])
+        x_test[..., channel_idx:channel_idx+1] = scaler.transform(x_test[..., channel_idx:channel_idx+1])
+
+    y_day = add_window_y(time_in_day, config['lag'], config['horizon'])
+    y_week = add_window_y(day_in_week, config['lag'], config['horizon'])
+
+    del time_in_day, day_in_week
+    gc.collect()
+
+    y = np.concatenate([y, y_day, y_week], axis=-1)
+
+    del y_day, y_week
+    gc.collect()
+
+    # Split Y
+    if config['test_ratio'] > 1:
+        y_train, y_val, y_test = split_data_by_days(y, config['val_ratio'], config['test_ratio'])
+    else:
+        y_train, y_val, y_test = split_data_by_ratio(y, config['val_ratio'], config['test_ratio'])
+
+    del y
+    gc.collect()
+
+    # Step 5: x_train y_train x_val y_val x_test y_test --> train val test
+    # train_dataloader = data_loader(x_train[..., :args['input_dim']], y_train[..., :args['input_dim']], args['batch_size'], shuffle=True, drop_last=True)
+    train_dataloader = data_loader(x_train, y_train, config['batch_size'], shuffle=True, drop_last=True)
+
+    del x_train, y_train
+    gc.collect()
+
+    # val_dataloader = data_loader(x_val[..., :args['input_dim']], y_val[..., :args['input_dim']], args['batch_size'], shuffle=False, drop_last=True)
+    val_dataloader = data_loader(x_val, y_val, config['batch_size'], shuffle=False, drop_last=True)
+
+    del x_val, y_val
+    gc.collect()
+
+    # test_dataloader = data_loader(x_test[..., :args['input_dim']], y_test[..., :args['input_dim']], args['batch_size'], shuffle=False, drop_last=False)
+    test_dataloader = data_loader(x_test, y_test, config['batch_size'], shuffle=False, drop_last=False)
+
+    del x_test, y_test
+    gc.collect()
+
+    return train_dataloader, val_dataloader, test_dataloader, scalers
+
+def split_data_by_days(data, val_days, test_days, interval=30):
+    t = int((24 * 60) / interval)
+    test_data = data[-t * int(test_days):]
+    val_data = data[-t * int(test_days + val_days):-t * int(test_days)]
+    train_data = data[:-t * int(test_days + val_days)]
+    return train_data, val_data, test_data
+
+
+def split_data_by_ratio(data, val_ratio, test_ratio):
+    data_len = data.shape[0]
+    test_data = data[-int(data_len * test_ratio):]
+    val_data = data[-int(data_len * (test_ratio + val_ratio)):-int(data_len * test_ratio)]
+    train_data = data[:-int(data_len * (test_ratio + val_ratio))]
+    return train_data, val_data, test_data
+
+
+def data_loader(X, Y, batch_size, shuffle=True, drop_last=True):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    X = torch.tensor(X, dtype=torch.float32, device=device)
+    Y = torch.tensor(Y, dtype=torch.float32, device=device)
+    data = torch.utils.data.TensorDataset(X, Y)
+    dataloader = torch.utils.data.DataLoader(data, batch_size=batch_size,
+                                             shuffle=shuffle, drop_last=drop_last)
+    return dataloader
+
+
+def add_window_x(data, window=3, horizon=1, single=False):
+    """
+    Generate windowed X values from the input data.
+
+    :param data: Input data, shape [B, ...]
+    :param window: Size of the sliding window
+    :param horizon: Horizon size
+    :param single: If True, generate single-step windows, else multi-step
+    :return: X with shape [B, W, ...]
+    """
+    length = len(data)
+    end_index = length - horizon - window + 1
+    x = []  # Sliding windows
+    index = 0
+
+    while index < end_index:
+        x.append(data[index:index + window])
+        index += 1
+
+    return np.array(x)
+
+
+def add_window_y(data, window=3, horizon=1, single=False):
+    """
+    Generate windowed Y values from the input data.
+
+    :param data: Input data, shape [B, ...]
+    :param horizon: Horizon size
+    :param single: If True, generate single-step windows, else multi-step
+    :return: Y with shape [B, H, ...]
+    """
+    length = len(data)
+    end_index = length - horizon - window + 1
+    y = []  # Horizon values
+    index = 0
+
+    while index < end_index:
+        if single:
+            y.append(data[index + window + horizon - 1:index + window + horizon])
+        else:
+            y.append(data[index + window:index + window + horizon])
+        index += 1
+
+    return np.array(y)

main.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+"""
+时空数据深度学习预测项目主程序
+专门处理时空数据格式 (batch_size, seq_len, num_nodes, features)
+"""
+
+import os
+from utils.args_reader import config_loader
+import utils.init as init
+import torch
+
+
+def main():
+    config = config_loader()
+    device = config['basic']['device'] = init.device(config['basic']['device'])
+    init.seed(config['basic']['seed'])
+    model = init.model(config)
+    train_loader, val_loader, test_loader, scaler = init.dataloader(config)
+    loss = init.loss(config, scaler)
+    optim, lr = init.optimizer(config, model)
+    logger = init.Logger(config)
+    trainer = init.trainer(config, model, loss, optim, train_loader, val_loader, test_loader, scaler, logger, lr)
+    match config['basic']['mode']:
+        case 'train':
+            trainer.train()
+        case 'test':
+            params_path = f"./pre-trained/{config['basic']['model']}/{config['basic']['dataset']}.pth"
+            params = torch.load(params_path, map_location=device, weights_only=True)
+            model.load_state_dict(params)
+            trainer.test(model.to(device), config, test_loader, scaler, trainer.logger)
+
+
+if __name__ == "__main__":
+    main()
+

models/model_selector.py

@@ -0,0 +1,6 @@
+
+
+def model_selector(config):
+    model_name = config['basic']['model']
+    model = None
+    return model

trainer/trainer.py

@@ -0,0 +1,519 @@
+import math
+import os
+import time
+import copy
+from tqdm import tqdm
+
+import torch
+
+class Trainer:
+    def __init__(self, config, model, loss, optimizer, train_loader, val_loader, test_loader,
+                 scalers, logger, lr_scheduler=None):
+        self.model = model
+        self.loss = loss
+        self.optimizer = optimizer
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.test_loader = test_loader
+        self.scalers = scalers  # 现在是多个标准化器的列表
+        self.args = config['train']
+        self.logger = logger
+        self.args['device'] = config['basic']['device']
+        self.lr_scheduler = lr_scheduler
+        self.train_per_epoch = len(train_loader)
+        self.val_per_epoch = len(val_loader) if val_loader else 0
+        self.best_path = os.path.join(logger.dir_path, 'best_model.pth')
+        self.best_test_path = os.path.join(logger.dir_path, 'best_test_model.pth')
+        self.loss_figure_path = os.path.join(logger.dir_path, 'loss.png')
+
+    def _run_epoch(self, epoch, dataloader, mode):
+        if mode == 'train':
+            self.model.train()
+            optimizer_step = True
+        else:
+            self.model.eval()
+            optimizer_step = False
+
+        total_loss = 0
+        epoch_time = time.time()
+
+        with torch.set_grad_enabled(optimizer_step):
+            with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar:
+                for batch_idx, (data, target) in enumerate(dataloader):
+                    label = target[..., :self.args['output_dim']]
+                    output = self.model(data).to(self.args['device'])
+
+                    if self.args['real_value']:
+                        # 只对输出维度进行反归一化
+                        output = self._inverse_transform_output(output)
+
+                    loss = self.loss(output, label)
+                    if optimizer_step and self.optimizer is not None:
+                        self.optimizer.zero_grad()
+                        loss.backward()
+
+                        if self.args['grad_norm']:
+                            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm'])
+                        self.optimizer.step()
+
+                    total_loss += loss.item()
+
+                    if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0:
+                        self.logger.info(
+                            f'Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {loss.item():.6f}')
+
+                    # 更新 tqdm 的进度
+                    pbar.update(1)
+                    pbar.set_postfix(loss=loss.item())
+
+        avg_loss = total_loss / len(dataloader)
+        self.logger.logger.info(
+            f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s')
+        return avg_loss
+
+    def _inverse_transform_output(self, output):
+        """
+        只对输出维度进行反归一化
+        假设输出数据形状为 [batch, horizon, nodes, features]
+        只对前output_dim个特征进行反归一化
+        """
+        if not self.args['real_value']:
+            return output
+            
+        # 获取输出维度的数量
+        output_dim = self.args['output_dim']
+        
+        # 如果输出特征数小于等于标准化器数量，直接使用对应的标准化器
+        if output_dim <= len(self.scalers):
+            # 对每个输出特征分别进行反归一化
+            for feature_idx in range(output_dim):
+                if feature_idx < len(self.scalers):
+                    output[..., feature_idx:feature_idx+1] = self.scalers[feature_idx].inverse_transform(
+                        output[..., feature_idx:feature_idx+1]
+                    )
+        else:
+            # 如果输出特征数大于标准化器数量，只对前len(scalers)个特征进行反归一化
+            for feature_idx in range(len(self.scalers)):
+                output[..., feature_idx:feature_idx+1] = self.scalers[feature_idx].inverse_transform(
+                    output[..., feature_idx:feature_idx+1]
+                )
+        
+        return output
+
+    def train_epoch(self, epoch):
+        return self._run_epoch(epoch, self.train_loader, 'train')
+
+    def val_epoch(self, epoch):
+        return self._run_epoch(epoch, self.val_loader or self.test_loader, 'val')
+
+    def test_epoch(self, epoch):
+        return self._run_epoch(epoch, self.test_loader, 'test')
+
+    def train(self):
+        best_model, best_test_model = None, None
+        best_loss, best_test_loss = float('inf'), float('inf')
+        not_improved_count = 0
+
+        self.logger.logger.info("Training process started")
+        for epoch in range(1, self.args['epochs'] + 1):
+            train_epoch_loss = self.train_epoch(epoch)
+            val_epoch_loss = self.val_epoch(epoch)
+            test_epoch_loss = self.test_epoch(epoch)
+
+            if train_epoch_loss > 1e6:
+                self.logger.logger.warning('Gradient explosion detected. Ending...')
+                break
+
+            if val_epoch_loss < best_loss:
+                best_loss = val_epoch_loss
+                not_improved_count = 0
+                best_model = copy.deepcopy(self.model.state_dict())
+                torch.save(best_model, self.best_path)
+                self.logger.logger.info('Best validation model saved!')
+            else:
+                not_improved_count += 1
+
+            if self.args['early_stop'] and not_improved_count == self.args['early_stop_patience']:
+                self.logger.logger.info(
+                    f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops.")
+                break
+
+            if test_epoch_loss < best_test_loss:
+                best_test_loss = test_epoch_loss
+                best_test_model = copy.deepcopy(self.model.state_dict())
+                torch.save(best_test_model, self.best_test_path)
+
+        if not self.args['debug']:
+            torch.save(best_model, self.best_path)
+            torch.save(best_test_model, self.best_test_path)
+            self.logger.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}")
+
+        self._finalize_training(best_model, best_test_model)
+
+    def _finalize_training(self, best_model, best_test_model):
+        self.model.load_state_dict(best_model)
+        self.logger.logger.info("Testing on best validation model")
+        self.test(self.model, self.args, self.test_loader, self.scalers, self.logger, generate_viz=False)
+
+        self.model.load_state_dict(best_test_model)
+        self.logger.logger.info("Testing on best test model")
+        self.test(self.model, self.args, self.test_loader, self.scalers, self.logger, generate_viz=True)
+
+    @staticmethod
+    def test(model, args, data_loader, scalers, logger, path=None, generate_viz=True):
+        if path:
+            checkpoint = torch.load(path)
+            model.load_state_dict(checkpoint['state_dict'])
+            model.to(args.device)
+
+        model.eval()
+        y_pred, y_true = [], []
+
+        with torch.no_grad():
+            for data, target in data_loader:
+                label = target[..., :args['output_dim']]
+                output = model(data)
+                y_pred.append(output)
+                y_true.append(label)
+
+        if args['real_value']:
+            # 只对输出维度进行反归一化
+            y_pred = Trainer._inverse_transform_output_static(torch.cat(y_pred, dim=0), args, scalers)
+        else:
+            y_pred = torch.cat(y_pred, dim=0)
+        y_true = torch.cat(y_true, dim=0)
+
+        # 计算每个时间步的指标
+        for t in range(y_true.shape[1]):
+            mae, rmse, mape = logger.all_metrics(y_pred[:, t, ...], y_true[:, t, ...],
+                                                 args['mae_thresh'], args['mape_thresh'])
+            logger.logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
+
+        mae, rmse, mape = logger.all_metrics(y_pred, y_true, args['mae_thresh'], args['mape_thresh'])
+        logger.logger.info(f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
+
+        # 只在需要时生成可视化图片
+        if generate_viz:
+            save_dir = logger.dir_path if hasattr(logger, 'dir_path') else './logs'
+            Trainer._generate_node_visualizations(y_pred, y_true, logger, save_dir)
+            Trainer._generate_input_output_comparison(y_pred, y_true, data_loader, logger, save_dir, 
+                                                   target_node=1, num_samples=10, scalers=scalers)
+
+    @staticmethod
+    def _inverse_transform_output_static(output, args, scalers):
+        """
+        静态方法：只对输出维度进行反归一化
+        """
+        if not args['real_value']:
+            return output
+            
+        # 获取输出维度的数量
+        output_dim = args['output_dim']
+        
+        # 如果输出特征数小于等于标准化器数量，直接使用对应的标准化器
+        if output_dim <= len(scalers):
+            # 对每个输出特征分别进行反归一化
+            for feature_idx in range(output_dim):
+                if feature_idx < len(scalers):
+                    output[..., feature_idx:feature_idx+1] = scalers[feature_idx].inverse_transform(
+                        output[..., feature_idx:feature_idx+1]
+                    )
+        else:
+            # 如果输出特征数大于标准化器数量，只对前len(scalers)个特征进行反归一化
+            for feature_idx in range(len(scalers)):
+                output[..., feature_idx:feature_idx+1] = scalers[feature_idx].inverse_transform(
+                    output[..., feature_idx:feature_idx+1]
+                )
+        
+        return output
+
+    @staticmethod
+    def _generate_node_visualizations(y_pred, y_true, logger, save_dir):
+        """
+        生成节点预测可视化图片
+
+        Args:
+            y_pred: 预测值
+            y_true: 真实值
+            logger: 日志记录器
+            save_dir: 保存目录
+        """
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import os
+        import matplotlib
+        from tqdm import tqdm
+
+        # 设置matplotlib配置，减少字体查找输出
+        matplotlib.set_loglevel('error')  # 只显示错误信息
+        plt.rcParams['font.family'] = 'DejaVu Sans'  # 使用默认字体
+
+        # 检查数据有效性
+        if y_pred is None or y_true is None:
+            return
+
+        # 创建pic文件夹
+        pic_dir = os.path.join(save_dir, 'pic')
+        os.makedirs(pic_dir, exist_ok=True)
+
+        # 固定生成10张图片
+        num_nodes_to_plot = 10
+
+        # 生成单个节点的详细图
+        with tqdm(total=num_nodes_to_plot, desc="Generating node visualizations") as pbar:
+            for node_id in range(num_nodes_to_plot):
+                # 获取对应节点的数据
+                if len(y_pred.shape) > 2 and y_pred.shape[-2] > node_id:
+                    # 数据格式: [time_step, seq_len, num_node, dim]
+                    node_pred = y_pred[:, 12, node_id, 0].cpu().numpy()  # t=1时刻，指定节点，第一个特征
+                    node_true = y_true[:, 12, node_id, 0].cpu().numpy()
+                else:
+                    # 如果数据不足10个节点，只处理实际存在的节点
+                    if node_id >= y_pred.shape[-2]:
+                        pbar.update(1)
+                        continue
+                    else:
+                        node_pred = y_pred[:, 0, node_id, 0].cpu().numpy()
+                        node_true = y_true[:, 0, node_id, 0].cpu().numpy()
+
+                # 检查数据有效性
+                if np.isnan(node_pred).any() or np.isnan(node_true).any():
+                    pbar.update(1)
+                    continue
+
+                # 取前500个时间步
+                max_steps = min(500, len(node_pred))
+                if max_steps <= 0:
+                    pbar.update(1)
+                    continue
+
+                node_pred_500 = node_pred[:max_steps]
+                node_true_500 = node_true[:max_steps]
+
+                # 创建时间轴
+                time_steps = np.arange(max_steps)
+
+                # 绘制对比图
+                plt.figure(figsize=(12, 6))
+                plt.plot(time_steps, node_true_500, 'b-', label='True Values', linewidth=2, alpha=0.8)
+                plt.plot(time_steps, node_pred_500, 'r-', label='Predictions', linewidth=2, alpha=0.8)
+                plt.xlabel('Time Steps')
+                plt.ylabel('Values')
+                plt.title(f'Node {node_id + 1}: True vs Predicted Values (First {max_steps} Time Steps)')
+                plt.legend()
+                plt.grid(True, alpha=0.3)
+
+                # 保存图片，使用不同的命名
+                save_path = os.path.join(pic_dir, f'node{node_id + 1:02d}_prediction_first500.png')
+                plt.savefig(save_path, dpi=300, bbox_inches='tight')
+                plt.close()
+
+                pbar.update(1)
+
+        # 生成所有节点的对比图（前100个时间步，便于观察）
+        # 选择前100个时间步
+        plot_steps = min(100, y_pred.shape[0])
+        if plot_steps <= 0:
+            return
+
+        # 创建子图
+        fig, axes = plt.subplots(2, 5, figsize=(20, 8))
+        axes = axes.flatten()
+
+        for node_id in range(num_nodes_to_plot):
+            if len(y_pred.shape) > 2 and y_pred.shape[-2] > node_id:
+                # 数据格式: [time_step, seq_len, num_node, dim]
+                node_pred = y_pred[:plot_steps, 0, node_id, 0].cpu().numpy()
+                node_true = y_true[:plot_steps, 0, node_id, 0].cpu().numpy()
+            else:
+                # 如果数据不足10个节点，只处理实际存在的节点
+                if node_id >= y_pred.shape[-2]:
+                    axes[node_id].text(0.5, 0.5, f'Node {node_id + 1}\nNo Data',
+                                       ha='center', va='center', transform=axes[node_id].transAxes)
+                    continue
+                else:
+                    node_pred = y_pred[:plot_steps, 0, node_id, 0].cpu().numpy()
+                    node_true = y_true[:plot_steps, 0, node_id, 0].cpu().numpy()
+
+            # 检查数据有效性
+            if np.isnan(node_pred).any() or np.isnan(node_true).any():
+                axes[node_id].text(0.5, 0.5, f'Node {node_id + 1}\nNo Data',
+                                   ha='center', va='center', transform=axes[node_id].transAxes)
+                continue
+
+            time_steps = np.arange(plot_steps)
+
+            axes[node_id].plot(time_steps, node_true, 'b-', label='True', linewidth=1.5, alpha=0.8)
+            axes[node_id].plot(time_steps, node_pred, 'r-', label='Pred', linewidth=1.5, alpha=0.8)
+            axes[node_id].set_title(f'Node {node_id + 1}')
+            axes[node_id].grid(True, alpha=0.3)
+            axes[node_id].legend(fontsize=8)
+
+            if node_id >= 5:  # 下面一行添加x轴标签
+                axes[node_id].set_xlabel('Time Steps')
+            if node_id % 5 == 0:  # 左边一列添加y轴标签
+                axes[node_id].set_ylabel('Values')
+
+        plt.tight_layout()
+        summary_path = os.path.join(pic_dir, 'all_nodes_summary.png')
+        plt.savefig(summary_path, dpi=300, bbox_inches='tight')
+        plt.close()
+
+    @staticmethod
+    def _generate_input_output_comparison(y_pred, y_true, data_loader, logger, save_dir, 
+                                        target_node=1, num_samples=10, scalers=None):
+        """
+        生成输入-输出样本比较图
+        
+        Args:
+            y_pred: 预测值
+            y_true: 真实值
+            data_loader: 数据加载器，用于获取输入数据
+            logger: 日志记录器
+            save_dir: 保存目录
+            target_node: 目标节点ID（从1开始）
+            num_samples: 要比较的样本数量
+            scalers: 标准化器列表，用于反归一化输入数据
+        """
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import os
+        import matplotlib
+        from tqdm import tqdm
+
+        # 设置matplotlib配置
+        matplotlib.set_loglevel('error')
+        plt.rcParams['font.family'] = 'DejaVu Sans'
+
+        # 创建compare文件夹
+        compare_dir = os.path.join(save_dir, 'pic', 'compare')
+        os.makedirs(compare_dir, exist_ok=True)
+
+        # 获取输入数据
+        input_data = []
+        for batch_idx, (data, target) in enumerate(data_loader):
+            if batch_idx >= num_samples:
+                break
+            input_data.append(data.cpu().numpy())
+        
+        if not input_data:
+            return
+
+        # 获取目标节点的索引（从0开始）
+        node_idx = target_node - 1
+        
+        # 检查节点索引是否有效
+        if node_idx >= y_pred.shape[-2]:
+            return
+
+        # 为每个样本生成比较图
+        with tqdm(total=min(num_samples, len(input_data)), desc="Generating input-output comparisons") as pbar:
+            for sample_idx in range(min(num_samples, len(input_data))):
+                # 获取输入序列（假设输入形状为 [batch, seq_len, nodes, features]）
+                input_seq = input_data[sample_idx][0, :, node_idx, 0]  # 第一个batch，所有时间步，目标节点，第一个特征
+                
+                # 对输入数据进行反归一化
+                if scalers is not None and len(scalers) > 0:
+                    # 使用第一个标准化器对输入进行反归一化（假设输入特征使用第一个标准化器）
+                    input_seq = scalers[0].inverse_transform(input_seq.reshape(-1, 1)).flatten()
+                
+                # 获取对应的预测值和真实值
+                pred_seq = y_pred[sample_idx, :, node_idx, 0].cpu().numpy()  # 所有horizon，目标节点，第一个特征
+                true_seq = y_true[sample_idx, :, node_idx, 0].cpu().numpy()
+                
+                # 检查数据有效性
+                if (np.isnan(input_seq).any() or np.isnan(pred_seq).any() or np.isnan(true_seq).any()):
+                    pbar.update(1)
+                    continue
+                
+                # 创建时间轴 - 输入和输出连续
+                total_time = np.arange(len(input_seq) + len(pred_seq))
+                
+                # 创建合并的图形 - 输入和输出在同一个图中
+                plt.figure(figsize=(14, 8))
+                
+                # 绘制完整的真实值曲线（输入 + 真实输出）
+                true_combined = np.concatenate([input_seq, true_seq])
+                plt.plot(total_time, true_combined, 'b', label='True Values (Input + Output)', 
+                        linewidth=2.5, alpha=0.9, linestyle='-')
+                
+                # 绘制预测值曲线（只绘制输出部分）
+                output_time = np.arange(len(input_seq), len(input_seq) + len(pred_seq))
+                plt.plot(output_time, pred_seq, 'r', label='Predicted Values', 
+                        linewidth=2, alpha=0.8, linestyle='-')
+                
+                # 添加垂直线分隔输入和输出
+                plt.axvline(x=len(input_seq)-0.5, color='gray', linestyle=':', alpha=0.7, 
+                           label='Input/Output Boundary')
+                
+                # 设置图形属性
+                plt.xlabel('Time Steps')
+                plt.ylabel('Values')
+                plt.title(f'Sample {sample_idx + 1}: Input-Output Comparison (Node {target_node})')
+                plt.legend()
+                plt.grid(True, alpha=0.3)
+                
+                # 调整布局
+                plt.tight_layout()
+                
+                # 保存图片
+                save_path = os.path.join(compare_dir, f'sample{sample_idx + 1:02d}_node{target_node:02d}_comparison.png')
+                plt.savefig(save_path, dpi=300, bbox_inches='tight')
+                plt.close()
+                
+                pbar.update(1)
+        
+        # 生成汇总图（所有样本的预测值对比）
+        
+        fig, axes = plt.subplots(2, 5, figsize=(20, 8))
+        axes = axes.flatten()
+        
+        for sample_idx in range(min(num_samples, len(input_data))):
+            if sample_idx >= 10:  # 最多显示10个子图
+                break
+                
+            ax = axes[sample_idx]
+            
+            # 获取输入序列和预测值、真实值
+            input_seq = input_data[sample_idx][0, :, node_idx, 0]
+            if scalers is not None and len(scalers) > 0:
+                input_seq = scalers[0].inverse_transform(input_seq.reshape(-1, 1)).flatten()
+            
+            pred_seq = y_pred[sample_idx, :, node_idx, 0].cpu().numpy()
+            true_seq = y_true[sample_idx, :, node_idx, 0].cpu().numpy()
+            
+            # 检查数据有效性
+            if np.isnan(input_seq).any() or np.isnan(pred_seq).any() or np.isnan(true_seq).any():
+                ax.text(0.5, 0.5, f'Sample {sample_idx + 1}\nNo Data',
+                       ha='center', va='center', transform=ax.transAxes)
+                continue
+            
+            # 绘制对比图 - 输入和输出连续显示
+            total_time = np.arange(len(input_seq) + len(pred_seq))
+            true_combined = np.concatenate([input_seq, true_seq])
+            output_time = np.arange(len(input_seq), len(input_seq) + len(pred_seq))
+            
+            ax.plot(total_time, true_combined, 'b', label='True', linewidth=2, alpha=0.9, linestyle='-')
+            ax.plot(output_time, pred_seq, 'r', label='Pred', linewidth=1.5, alpha=0.8, linestyle='-')
+            ax.axvline(x=len(input_seq)-0.5, color='gray', linestyle=':', alpha=0.5)
+            ax.set_title(f'Sample {sample_idx + 1}')
+            ax.grid(True, alpha=0.3)
+            ax.legend(fontsize=8)
+            
+            if sample_idx >= 5:  # 下面一行添加x轴标签
+                ax.set_xlabel('Time Steps')
+            if sample_idx % 5 == 0:  # 左边一列添加y轴标签
+                ax.set_ylabel('Values')
+        
+        # 隐藏多余的子图
+        for i in range(min(num_samples, len(input_data)), 10):
+            axes[i].set_visible(False)
+        
+        plt.tight_layout()
+        summary_path = os.path.join(compare_dir, f'all_samples_node{target_node:02d}_summary.png')
+        plt.savefig(summary_path, dpi=300, bbox_inches='tight')
+        plt.close()
+
+    @staticmethod
+    def _compute_sampling_threshold(global_step, k):
+        return k / (k + math.exp(global_step / k))

trainer/trainer_selector.py

@@ -0,0 +1,11 @@
+from trainer.trainer import Trainer
+
+def select_trainer(config, model, loss, optimizer, train_loader, val_loader, test_loader, scaler,
+                   lr_scheduler, kwargs):
+    model_name = config['basic']['model']
+    selected_Trainer = None
+    match model_name:
+        case _: selected_Trainer = Trainer(config, model, loss, optimizer,
+                               train_loader, val_loader, test_loader, scaler,lr_scheduler)
+    if selected_Trainer is None: raise NotImplementedError
+    return selected_Trainer

utils/args_reader.py

@@ -0,0 +1,12 @@
+import argparse
+import yaml 
+
+def config_loader():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="./config/DDGCRN_config.yaml")
+    
+    config_path = parser.parse_args().config
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+

utils/init.py

@@ -0,0 +1,173 @@
+import os
+import torch
+import torch.nn as nn
+import random
+import yaml
+import logging
+from datetime import datetime
+import numpy as np
+
+from models.model_selector import model_selector
+from data.data_selector import load_dataset
+from data.dataloader import get_dataloader
+import utils.loss_func as loss_func
+from trainer.trainer_selector import select_trainer
+
+
+def seed(seed : int):
+    """ 固定随机种子以公平测试 """
+    torch.cuda.cudnn_enabled = False
+    torch.backends.cudnn.deterministic = True
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    # print(f"seed is {seed}")
+
+def device(device : str):
+    """初始化使用设备"""
+    if torch.cuda.is_available() and device != 'cpu':
+        torch.cuda.set_device(int(device.split(':')[1]))
+        return device
+    else:
+        return 'cpu'
+
+def model(config : dict):
+    """选择模型"""
+    device = config['basic']['device']
+    model = model_selector(config).to(device)
+    for p in model.parameters():
+        if p.dim() > 1: nn.init.xavier_uniform_(p)
+        else: nn.init.uniform_(p)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model param count : {total_params}")
+    return model
+
+def dataloader(config : dict):
+    """初始化dataloader"""
+    data = load_dataset(config)
+    train_loader, val_loader, test_loader, scaler = get_dataloader(config, data)
+    return train_loader, val_loader, test_loader, scaler
+
+def loss(config : dict, scaler):
+    loss_name = config['train']['loss']
+    device = config['basic']['device']
+    match loss_name :
+        case 'mask_mae': func = loss_func.masked_mae_loss(scaler, mask_value=0.0)
+        case 'mae': func = torch.nn.L1Loss()
+        case 'mse': func = torch.nn.MSELoss()
+        case 'Huber': func = torch.nn.HuberLoss()
+        case _ : raise NotImplementedError('No Loss Func')
+    return func.to(device)
+
+
+def optimizer(config, model):
+    optimizer = torch.optim.Adam(
+        params=model.parameters(),
+        lr=config['train']['lr_init'],
+        eps=1.0e-8,
+        weight_decay=config['train']['weight_decay'],
+        amsgrad=False
+    )
+
+    lr_scheduler = None
+    if config['train']['lr_decay']:
+        lr_decay_steps = [int(step) for step in config['train']['lr_decay_step'].split(',')]
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer=optimizer,
+            milestones=lr_decay_steps,
+            gamma=config['train']['lr_decay_rate']
+        )
+
+    return optimizer, lr_scheduler
+
+
+def trainer(config, model, loss, optimizer,
+            train_loader, val_loader, test_loader,
+            scaler, lr_scheduler, kwargs):
+    selected_trainer = select_trainer(config, model, loss, optimizer,
+            train_loader, val_loader, test_loader, scaler, lr_scheduler, kwargs)
+    return selected_trainer
+
+class Logger:
+    """
+    Logger类，主要调用成员对象logger的info方法来记录
+    使用logger的all_metrics返回所有损失
+    """
+    def __init__(self, config, name=None, debug = True):
+        self.config = config
+        cur_time = datetime.now().strftime("%Y/%m/%d-%H:%M:%S")
+        cur_dir = os.getcwd()
+        dataset_name = config['basic']['dataset']
+        model_name = config['basic']['model']
+        self.dir_path = os.path.join(cur_dir, 'exp', f'{dataset_name}_{model_name}_{cur_time}')
+        config['train']['log_dir'] = self.dir_path
+        os.makedirs(self.dir_path, exist_ok=True)
+        # 生成配置并添加到目录
+        config_content = yaml.safe_dump(config)
+        config_path = os.path.join(self.dir_path, "config.yaml")
+        with open(config_path, 'w') as f:
+            f.write(config_content)
+
+        # logger
+        self.logger = logging.getLogger(name)
+        self.logger.setLevel(logging.DEBUG)
+        formatter = logging.Formatter('%(asctime)s: %(message)s', "%m/%d %H:%M")
+        
+        # 控制台处理器
+        console_handler = logging.StreamHandler()
+        if debug:
+            console_handler.setLevel(logging.DEBUG)
+        else:
+            console_handler.setLevel(logging.INFO)
+        console_handler.setFormatter(formatter)
+        
+        # 文件处理器 - 无论是否debug都创建日志文件
+        logfile = os.path.join(self.dir_path, 'run.log')
+        file_handler = logging.FileHandler(logfile, mode='w')
+        file_handler.setLevel(logging.DEBUG)
+        file_handler.setFormatter(formatter)
+        
+        # 添加处理器到logger
+        self.logger.addHandler(console_handler)
+        self.logger.addHandler(file_handler)
+
+    def set_log_dir(self):
+        # Initialize logger
+        if not os.path.isdir(self.dir_path) and not self.config['basic']['debug']:
+            os.makedirs(self.dir_path, exist_ok=True)
+        self.logger.info(f"Experiment log path in: {self.dir_path}")
+
+    def mae_torch(self, pred, true, mask_value=None):
+        if mask_value is not None:
+            mask = torch.gt(true, mask_value)
+            pred = torch.masked_select(pred, mask)
+            true = torch.masked_select(true, mask)
+        return torch.mean(torch.abs(true - pred))
+
+    def rmse_torch(self, pred, true, mask_value=None):
+        if mask_value is not None:
+            mask = torch.gt(true, mask_value)
+            pred = torch.masked_select(pred, mask)
+            true = torch.masked_select(true, mask)
+        return torch.sqrt(torch.mean((pred - true) ** 2))
+
+    def mape_torch(self, pred, true, mask_value=None):
+        if mask_value is not None:
+            mask = torch.gt(true, mask_value)
+            pred = torch.masked_select(pred, mask)
+            true = torch.masked_select(true, mask)
+        return torch.mean(torch.abs(torch.div((true - pred), (true + 0.001))))
+
+    def all_metrics(self, pred, true, mask1, mask2):
+        if mask1 == 'None': mask1 = None
+        if mask2 == 'None': mask2 = None
+        mae = self.mae_torch(pred, true, mask1)
+        rmse = self.rmse_torch(pred, true, mask1)
+        mape = self.mape_torch(pred, true, mask2)
+        return mae, rmse, mape
+
+
+
+
+

utils/loss_func.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+
+class MaskedMAELoss(nn.Module):
+    def __init__(self, scaler, mask_value):
+        super(MaskedMAELoss, self).__init__()
+        self.scaler = scaler
+        self.mask_value = mask_value
+    
+    def forward(self, preds, labels):
+        if self.scaler:
+            preds = self.scaler.inverse_transform(preds)
+            labels = self.scaler.inverse_transform(labels)
+        return mae_torch(pred=preds, true=labels, mask_value=self.mask_value)
+
+def masked_mae_loss(scaler, mask_value):
+    """保持向后兼容性的函数"""
+    return MaskedMAELoss(scaler, mask_value)
+
+def mae_torch(pred, true, mask_value=None):
+    if mask_value is not None:
+        mask = torch.gt(true, mask_value)
+        pred = torch.masked_select(pred, mask)
+        true = torch.masked_select(true, mask)
+    return torch.mean(torch.abs(true - pred))
+
+
+def rmse_torch(pred, true, mask_value=None):
+    if mask_value is not None:
+        mask = torch.gt(true, mask_value)
+        pred = torch.masked_select(pred, mask)
+        true = torch.masked_select(true, mask)
+    return torch.sqrt(torch.mean((pred - true) ** 2))
+
+
+def mape_torch(pred, true, mask_value=None):
+    if mask_value is not None:
+        mask = torch.gt(true, mask_value)
+        pred = torch.masked_select(pred, mask)
+        true = torch.masked_select(true, mask)
+    return torch.mean(torch.abs(torch.div((true - pred), (true + 0.001))))
+
+
+def all_metrics(pred, true, mask1, mask2):
+    if mask1 == 'None': mask1 = None
+    if mask2 == 'None': mask2 = None
+    mae = mae_torch(pred, true, mask1)
+    rmse = rmse_torch(pred, true, mask1)
+    mape = mape_torch(pred, true, mask2)
+    return mae, rmse, mape
+
+
+if __name__ == '__main__':
+    pred = torch.Tensor([1, 2, 3, 4])
+    true = torch.Tensor([2, 1, 4, 5])
+    print(all_metrics(pred, true, None, None))

utils/metrics.py

@@ -0,0 +1,93 @@
+import numpy as np
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+from typing import Dict, Union
+
+
+def calculate_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
+    """
+    计算评估指标
+    
+    Args:
+        y_true: 真实值
+        y_pred: 预测值
+        
+    Returns:
+        包含各种指标的字典
+    """
+    # 确保输入是numpy数组
+    y_true = np.array(y_true)
+    y_pred = np.array(y_pred)
+    
+    # 计算各种指标
+    mse = mean_squared_error(y_true, y_pred)
+    rmse = np.sqrt(mse)
+    mae = mean_absolute_error(y_true, y_pred)
+    
+    # 计算MAPE
+    mape = np.mean(np.abs((y_true - y_pred) / (y_true + 1e-8))) * 100
+    
+    # 计算R²
+    ss_res = np.sum((y_true - y_pred) ** 2)
+    ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)
+    r2 = 1 - (ss_res / (ss_tot + 1e-8))
+    
+    # 计算SMAPE
+    smape = 2.0 * np.mean(np.abs(y_pred - y_true) / (np.abs(y_true) + np.abs(y_pred) + 1e-8)) * 100
+    
+    metrics = {
+        'MSE': mse,
+        'RMSE': rmse,
+        'MAE': mae,
+        'MAPE': mape,
+        'R2': r2,
+        'SMAPE': smape
+    }
+    
+    return metrics
+
+
+def calculate_rolling_metrics(y_true: np.ndarray, y_pred: np.ndarray, 
+                            window: int = 10) -> Dict[str, np.ndarray]:
+    """
+    计算滚动评估指标
+    
+    Args:
+        y_true: 真实值
+        y_pred: 预测值
+        window: 滚动窗口大小
+        
+    Returns:
+        包含滚动指标的字典
+    """
+    y_true = np.array(y_true)
+    y_pred = np.array(y_pred)
+    
+    n = len(y_true)
+    if n < window:
+        return {}
+    
+    rolling_mse = []
+    rolling_mae = []
+    rolling_mape = []
+    
+    for i in range(window, n + 1):
+        start_idx = i - window
+        end_idx = i
+        
+        true_window = y_true[start_idx:end_idx]
+        pred_window = y_pred[start_idx:end_idx]
+        
+        # 计算窗口内的指标
+        mse = mean_squared_error(true_window, pred_window)
+        mae = mean_absolute_error(true_window, pred_window)
+        mape = np.mean(np.abs((true_window - pred_window) / (true_window + 1e-8))) * 100
+        
+        rolling_mse.append(mse)
+        rolling_mae.append(mae)
+        rolling_mape.append(mape)
+    
+    return {
+        'rolling_MSE': np.array(rolling_mse),
+        'rolling_MAE': np.array(rolling_mae),
+        'rolling_MAPE': np.array(rolling_mape)
+    }

utils/normalizer.py

@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+
+
+class NScaler:
+    """No normalization, returns the data as is."""
+
+    def transform(self, data):
+        return data
+
+    def inverse_transform(self, data):
+        return data
+
+
+class StandardScaler:
+    """Standardizes the input data by removing the mean and scaling to unit variance."""
+
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def transform(self, data):
+        return (data - self.mean) / self.std
+
+    def inverse_transform(self, data):
+        if isinstance(data, torch.Tensor) and isinstance(self.mean, np.ndarray):
+            self.std = torch.from_numpy(self.std).to(data.device).type(data.dtype)
+            self.mean = torch.from_numpy(self.mean).to(data.device).type(data.dtype)
+        return (data * self.std) + self.mean
+
+
+class MinMax01Scaler:
+    """Scales data to the range [0, 1]."""
+
+    def __init__(self, min, max):
+        self.min = min
+        self.max = max
+
+    def transform(self, data):
+        return (data - self.min) / (self.max - self.min)
+
+    def inverse_transform(self, data):
+        if isinstance(data, torch.Tensor) and isinstance(self.min, np.ndarray):
+            self.min = torch.from_numpy(self.min).to(data.device).type(data.dtype)
+            self.max = torch.from_numpy(self.max).to(data.device).type(data.dtype)
+        return (data * (self.max - self.min)) + self.min
+
+
+class MinMax11Scaler:
+    """Scales data to the range [-1, 1]."""
+
+    def __init__(self, min, max):
+        self.min = min
+        self.max = max
+
+    def transform(self, data):
+        return ((data - self.min) / (self.max - self.min)) * 2.0 - 1.0
+
+    def inverse_transform(self, data):
+        if isinstance(data, torch.Tensor) and isinstance(self.min, np.ndarray):
+            self.min = torch.from_numpy(self.min).to(data.device).type(data.dtype)
+            self.max = torch.from_numpy(self.max).to(data.device).type(data.dtype)
+        return ((data + 1.0) / 2.0) * (self.max - self.min) + self.min
+
+
+class ColumnMinMaxScaler:
+    """Scales data using column-specific min and max values."""
+
+    def __init__(self, min, max):
+        self.min = min
+        self.min_max = max - self.min
+        self.min_max[self.min_max == 0] = 1
+
+    def transform(self, data):
+        return (data - self.min) / self.min_max
+
+    def inverse_transform(self, data):
+        if isinstance(data, torch.Tensor) and isinstance(self.min, np.ndarray):
+            self.min_max = torch.from_numpy(self.min_max).to(data.device).type(torch.float32)
+            self.min = torch.from_numpy(self.min).to(data.device).type(torch.float32)
+        return (data * self.min_max) + self.min
+
+
+def one_hot_by_column(data):
+    """Applies one-hot encoding to each column of a 2D numpy array."""
+    len_data = data.shape[0]
+    encoded = []
+
+    for i in range(data.shape[1]):
+        column = data[:, i]
+        min_val = column.min()
+        zero_matrix = np.zeros((len_data, column.max() - min_val + 1))
+        zero_matrix[np.arange(len_data), column - min_val] = 1
+        encoded.append(zero_matrix)
+
+    return np.hstack(encoded)
+
+
+def minmax_by_column(data):
+    """Applies MinMax scaling to each column of a 2D numpy array."""
+    normalized = []
+
+    for i in range(data.shape[1]):
+        column = data[:, i]
+        min_val = column.min()
+        max_val = column.max()
+        column = (column - min_val) / (max_val - min_val)
+        normalized.append(column[:, np.newaxis])
+
+    return np.hstack(normalized)
+
+
+def normalize_dataset(data, normalizer, column_wise=False):
+    if normalizer == 'max01':
+        if column_wise:
+            minimum = data.min(axis=0, keepdims=True)
+            maximum = data.max(axis=0, keepdims=True)
+        else:
+            minimum = data.min()
+            maximum = data.max()
+        scaler = MinMax01Scaler(minimum, maximum)
+        # data = scaler.transform(data)
+        # print('Normalize the dataset by MinMax01 Normalization')
+    elif normalizer == 'max11':
+        if column_wise:
+            minimum = data.min(axis=0, keepdims=True)
+            maximum = data.max(axis=0, keepdims=True)
+        else:
+            minimum = data.min()
+            maximum = data.max()
+        scaler = MinMax11Scaler(minimum, maximum)
+        # data = scaler.transform(data)
+        # print('Normalize the dataset by MinMax11 Normalization')
+    elif normalizer == 'std':
+        if column_wise:
+            mean = data.mean(axis=0, keepdims=True)
+            std = data.std(axis=0, keepdims=True)
+        else:
+            mean = data.mean()
+            std = data.std()
+        scaler = StandardScaler(mean, std)
+        # data = scaler.transform(data)
+        # print('Normalize the dataset by Standard Normalization')
+    elif normalizer == 'None':
+        scaler = NScaler()
+        # data = scaler.transform(data)
+        # print('Does not normalize the dataset')
+    elif normalizer == 'cmax':
+        scaler = ColumnMinMaxScaler(data.min(axis=0), data.max(axis=0))
+        # data = scaler.transform(data)
+        # print('Normalize the dataset by Column Min-Max Normalization')
+    else:
+        raise ValueError(f"Unsupported normalizer type: {normalizer}")
+    return scaler