242 lines
9.3 KiB
Python
Executable File
242 lines
9.3 KiB
Python
Executable File
from lib.normalization import normalize_dataset
|
|
import numpy as np
|
|
import gc
|
|
import os
|
|
import torch
|
|
import h5py
|
|
|
|
from model.STGNCDE import controldiffeq
|
|
|
|
|
|
def get_dataloader(args, normalizer='std', single=True):
|
|
data = load_st_dataset(args['type']) # 加载数据
|
|
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()
|
|
|
|
data_category = 'traffic'
|
|
if data_category == 'traffic':
|
|
times = torch.linspace(0, 11, 12)
|
|
elif data_category == 'token':
|
|
times = torch.linspace(0, 6, 7)
|
|
else:
|
|
raise ValueError
|
|
augmented_X_tra = []
|
|
augmented_X_tra.append(
|
|
times.unsqueeze(0).unsqueeze(0).repeat(x_train.shape[0], x_train.shape[2], 1).unsqueeze(-1).transpose(1, 2))
|
|
augmented_X_tra.append(torch.Tensor(x_train[..., :]))
|
|
x_train = torch.cat(augmented_X_tra, dim=3)
|
|
augmented_X_val = []
|
|
augmented_X_val.append(
|
|
times.unsqueeze(0).unsqueeze(0).repeat(x_val.shape[0], x_val.shape[2], 1).unsqueeze(-1).transpose(1, 2))
|
|
augmented_X_val.append(torch.Tensor(x_val[..., :]))
|
|
x_val = torch.cat(augmented_X_val, dim=3)
|
|
augmented_X_test = []
|
|
augmented_X_test.append(
|
|
times.unsqueeze(0).unsqueeze(0).repeat(x_test.shape[0], x_test.shape[2], 1).unsqueeze(-1).transpose(1, 2))
|
|
augmented_X_test.append(torch.Tensor(x_test[..., :]))
|
|
x_test = torch.cat(augmented_X_test, dim=3)
|
|
|
|
train_coeffs = controldiffeq.natural_cubic_spline_coeffs(times, x_train.transpose(1, 2))
|
|
valid_coeffs = controldiffeq.natural_cubic_spline_coeffs(times, x_val.transpose(1, 2))
|
|
test_coeffs = controldiffeq.natural_cubic_spline_coeffs(times, x_test.transpose(1, 2))
|
|
|
|
train_dataloader = data_loader_cde(train_coeffs, y_train, args['batch_size'], shuffle=True, drop_last=True)
|
|
del x_train, y_train
|
|
gc.collect()
|
|
|
|
if len(x_val) == 0:
|
|
val_dataloader = None
|
|
else:
|
|
val_dataloader = data_loader_cde(valid_coeffs, y_val, args['batch_size'], shuffle=False, drop_last=True)
|
|
del x_val, y_val
|
|
gc.collect()
|
|
|
|
test_dataloader = data_loader_cde(test_coeffs, 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, times
|
|
|
|
|
|
def data_loader_cde(X, Y, batch_size, shuffle=True, drop_last=True):
|
|
cuda = True if torch.cuda.is_available() else False
|
|
TensorFloat = torch.cuda.FloatTensor if cuda else torch.FloatTensor
|
|
# X, Y = TensorFloat(X), TensorFloat(Y)
|
|
# X = tuple(TensorFloat(x) for x in X)
|
|
# Y = TensorFloat(Y)
|
|
data = torch.utils.data.TensorDataset(*X, torch.tensor(Y))
|
|
dataloader = torch.utils.data.DataLoader(data, batch_size=batch_size,
|
|
shuffle=shuffle, drop_last=drop_last)
|
|
return dataloader
|
|
|
|
|
|
def load_st_dataset(dataset):
|
|
# 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 _:
|
|
raise ValueError(f"Unsupported dataset: {dataset}")
|
|
|
|
# Ensure data shape compatibility
|
|
if len(data.shape) == 2:
|
|
data = np.expand_dims(data, axis=-1)
|
|
|
|
print('Load %s Dataset shaped: ' % dataset, data.shape, data.max(), data.min(), data.mean(), np.median(data))
|
|
return data
|
|
|
|
|
|
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)
|