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)