From 1be0b59344ad50eb8e238100eae2088790028ba2 Mon Sep 17 00:00:00 2001 From: czzhangheng Date: Thu, 8 May 2025 22:43:26 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9EEXP=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E5=8A=A0=E8=BD=BD=E5=99=A8=E3=80=81=E6=A8=A1=E5=9E=8B=E5=92=8C?= =?UTF-8?q?=E8=AE=AD=E7=BB=83=E5=99=A8=EF=BC=8C=E6=94=AF=E6=8C=81=E5=91=A8?= =?UTF-8?q?=E6=9C=9F=E6=80=A7=E6=95=B0=E6=8D=AE=E5=A4=84=E7=90=86=E5=92=8C?= =?UTF-8?q?=E5=8A=A8=E6=80=81=E5=9B=BE=E6=9E=84=E5=BB=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dataloader/EXPdataloader.py | 212 ++++++++++++++++++++++++++++++++++++ model/EXP/EXP32.py | 0 trainer/E32Trainer.py | 183 +++++++++++++++++++++++++++++++ 3 files changed, 395 insertions(+) create mode 100755 dataloader/EXPdataloader.py create mode 100644 model/EXP/EXP32.py create mode 100644 trainer/E32Trainer.py diff --git a/dataloader/EXPdataloader.py b/dataloader/EXPdataloader.py new file mode 100755 index 0000000..ad6d8e5 --- /dev/null +++ b/dataloader/EXPdataloader.py @@ -0,0 +1,212 @@ +from lib.normalization import normalize_dataset + +import numpy as np +import gc +import os +import torch +import h5py + + +def get_dataloader(args, normalizer='std', single=True): + data = load_st_dataset(args['type'], args['sample']) # 加载数据 + L, N, F = data.shape # 数据形状 + + # Step 1: data -> x,y + x = add_window_x(data, args['lag'], args['horizon'], single) + y = add_window_y(data, args['lag'], args['horizon'], single) + + del data + gc.collect() + + # Step 2: time_in_day, day_in_week -> day, week + time_in_day = [i % args['steps_per_day'] / args['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 // args['steps_per_day']) % args['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, args['lag'], args['horizon'], single) + x_week = add_window_x(day_in_week, args['lag'], args['horizon'], single) + + # 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 args['test_ratio'] > 1: + x_train, x_val, x_test = split_data_by_days(x, args['val_ratio'], args['test_ratio']) + else: + x_train, x_val, x_test = split_data_by_ratio(x, args['val_ratio'], args['test_ratio']) + + del x + gc.collect() + + # Normalization + scaler = normalize_dataset(x_train[..., :args['input_dim']], normalizer, args['column_wise']) + x_train[..., :args['input_dim']] = scaler.transform(x_train[..., :args['input_dim']]) + x_val[..., :args['input_dim']] = scaler.transform(x_val[..., :args['input_dim']]) + x_test[..., :args['input_dim']] = scaler.transform(x_test[..., :args['input_dim']]) + + + y_day = add_window_y(time_in_day, args['lag'], args['horizon'], single) + y_week = add_window_y(day_in_week, args['lag'], args['horizon'], single) + + 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 args['test_ratio'] > 1: + y_train, y_val, y_test = split_data_by_days(y, args['val_ratio'], args['test_ratio']) + else: + y_train, y_val, y_test = split_data_by_ratio(y, args['val_ratio'], args['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, args['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, args['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, args['batch_size'], shuffle=False, drop_last=False) + + del x_test, y_test + gc.collect() + + return train_dataloader, val_dataloader, test_dataloader, scaler + +def load_st_dataset(dataset, sample): + # output B, N, D + match dataset: + case 'PEMSD3': + data_path = os.path.join('./data/PEMS03/PEMS03.npz') + data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data + case 'PEMSD4': + data_path = os.path.join('./data/PEMS04/PEMS04.npz') + data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data + case 'PEMSD7': + data_path = os.path.join('./data/PEMS07/PEMS07.npz') + data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data + case 'PEMSD8': + data_path = os.path.join('./data/PEMS08/PEMS08.npz') + data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data + case 'PEMSD7(L)': + data_path = os.path.join('./data/PEMS07(L)/PEMS07L.npz') + data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data + case 'PEMSD7(M)': + data_path = os.path.join('./data/PEMS07(M)/V_228.csv') + data = np.genfromtxt(data_path, delimiter=',') # Read CSV directly with numpy + case 'METR-LA': + data_path = os.path.join('./data/METR-LA/METR.h5') + with h5py.File(data_path, 'r') as f: # Use h5py to handle HDF5 files without pandas + data = np.array(f['data']) + case 'BJ': + data_path = os.path.join('./data/BJ/BJ500.csv') + data = np.genfromtxt(data_path, delimiter=',', skip_header=1) # Skip header if present + case 'Hainan': + data_path = os.path.join('./data/Hainan/Hainan.npz') + data = np.load(data_path)['data'][:, :, 0] + case 'SD': + data_path = os.path.join('./data/SD/data.npz') + data = np.load(data_path)["data"][:, :, 0].astype(np.float32) + case _: + raise ValueError(f"Unsupported dataset: {dataset}") + + # Ensure data shape compatibility + if len(data.shape) == 2: + data = np.expand_dims(data, axis=-1) + + print('加载 %s 数据集中... ' % dataset) + return data[::sample] + +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 window: Size of the sliding window + :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) + + +if __name__ == '__main__': + res = load_st_dataset('SD', 1) + k = 1 \ No newline at end of file diff --git a/model/EXP/EXP32.py b/model/EXP/EXP32.py new file mode 100644 index 0000000..e69de29 diff --git a/trainer/E32Trainer.py b/trainer/E32Trainer.py new file mode 100644 index 0000000..8aa2bfe --- /dev/null +++ b/trainer/E32Trainer.py @@ -0,0 +1,183 @@ +import math +import os +import time +import copy +from tqdm import tqdm + +import torch +from lib.logger import get_logger +from lib.loss_function import all_metrics + + +class Trainer: + def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader, + scaler, args, 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.scaler = scaler + self.args = args + self.lr_scheduler = lr_scheduler + self.train_per_epoch = len(train_loader) + self.val_per_epoch = len(val_loader) if val_loader else 0 + + # Paths for saving models and logs + self.best_path = os.path.join(args['log_dir'], 'best_model.pth') + self.best_test_path = os.path.join(args['log_dir'], 'best_test_model.pth') + self.loss_figure_path = os.path.join(args['log_dir'], 'loss.png') + + # Initialize logger + if not os.path.isdir(args['log_dir']) and not args['debug']: + os.makedirs(args['log_dir'], exist_ok=True) + self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) + self.logger.info(f"Experiment log path in: {args['log_dir']}") + + def _run_epoch(self, epoch, dataloader, mode): + is_train = (mode == 'train') + self.model.train() if is_train else self.model.eval() + total_loss = 0.0 + + with torch.set_grad_enabled(is_train), \ + tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: + + for batch_idx, batch in enumerate(dataloader): + # unpack the new cycle_index + data, target, cycle_index = batch + data = data.to(self.args['device']) + target = target.to(self.args['device']) + cycle_index = cycle_index.to(self.args['device']).long() + + # forward + if is_train: + self.optimizer.zero_grad() + output = self.model(data, cycle_index) + else: + output = self.model(data, cycle_index) + + # compute loss + label = target[..., :self.args['output_dim']] + if self.args['real_value']: + output = self.scaler.inverse_transform(output) + loss = self.loss(output, label) + + # backward / step + if is_train: + 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() + + # logging + if is_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}' + ) + + pbar.update(1) + pbar.set_postfix(loss=loss.item()) + + avg_loss = total_loss / len(dataloader) + self.logger.info(f'{mode.capitalize()} Epoch {epoch}: avg Loss: {avg_loss:.6f}') + return avg_loss + + 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.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.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()) + self.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.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()) + + if not self.args['debug']: + torch.save(best_model, self.best_path) + torch.save(best_test_model, self.best_test_path) + self.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.info("Testing on best validation model") + self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) + + self.model.load_state_dict(best_test_model) + self.logger.info("Testing on best test model") + self.test(self.model, self.args, self.test_loader, self.scaler, self.logger) + + @staticmethod + def test(model, args, data_loader, scaler, logger, path=None): + 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 = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + else: + y_pred = torch.cat(y_pred, dim=0) + y_true = torch.cat(y_true, dim=0) + + # 你在这里需要把y_pred和y_true保存下来 + # torch.save(y_pred, "./test/PEMS07/y_pred_D.pt") # [3566,12,170,1] + # torch.save(y_true, "./test/PEMS08/y_true.pt") # [3566,12,170,1] + + for t in range(y_true.shape[1]): + mae, rmse, mape = all_metrics(y_pred[:, t, ...], y_true[:, t, ...], + args['mae_thresh'], args['mape_thresh']) + logger.info(f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") + + mae, rmse, mape = all_metrics(y_pred, y_true, args['mae_thresh'], args['mape_thresh']) + logger.info(f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") + + @staticmethod + def _compute_sampling_threshold(global_step, k): + return k / (k + math.exp(global_step / k))