新增EXP数据加载器、模型和训练器，支持周期性数据处理和动态图构建
This commit is contained in:
parent
eb8684bf91
commit
1be0b59344
|
|
@ -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
|
||||||
|
|
@ -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))
|
||||||
Loading…
Reference in New Issue