diff --git a/data/data_selector.py b/data/data_selector.py new file mode 100644 index 0000000..8ff652a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/data/dataloader.py b/data/dataloader.py new file mode 100644 index 0000000..ec7456d --- /dev/null +++ b/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) diff --git a/main.py b/main.py new file mode 100644 index 0000000..ea2e00d --- /dev/null +++ b/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() + diff --git a/models/model_selector.py b/models/model_selector.py new file mode 100644 index 0000000..fced561 --- /dev/null +++ b/models/model_selector.py @@ -0,0 +1,6 @@ + + +def model_selector(config): + model_name = config['basic']['model'] + model = None + return model \ No newline at end of file diff --git a/trainer/trainer.py b/trainer/trainer.py new file mode 100644 index 0000000..898cd84 --- /dev/null +++ b/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)) diff --git a/trainer/trainer_selector.py b/trainer/trainer_selector.py new file mode 100644 index 0000000..d5a3751 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/utils/args_reader.py b/utils/args_reader.py new file mode 100644 index 0000000..c786c84 --- /dev/null +++ b/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 + diff --git a/utils/init.py b/utils/init.py new file mode 100644 index 0000000..d8110db --- /dev/null +++ b/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 + + + + + diff --git a/utils/loss_func.py b/utils/loss_func.py new file mode 100644 index 0000000..6576803 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/utils/metrics.py b/utils/metrics.py new file mode 100644 index 0000000..95897f8 --- /dev/null +++ b/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) + } diff --git a/utils/normalizer.py b/utils/normalizer.py new file mode 100644 index 0000000..e269060 --- /dev/null +++ b/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 \ No newline at end of file