优化trainer和run
This commit is contained in:
parent
dceae4b1a3
commit
8b7e13df30
|
|
@ -1,5 +1,5 @@
|
|||
{
|
||||
"python-envs.defaultEnvManager": "ms-python.python:conda",
|
||||
"python-envs.defaultPackageManager": "ms-python.python:conda",
|
||||
"python-envs.defaultEnvManager": "ms-python.python:system",
|
||||
"python-envs.defaultPackageManager": "ms-python.python:pip",
|
||||
"python-envs.pythonProjects": []
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@ basic:
|
|||
mode : "train"
|
||||
device : "cuda:1"
|
||||
model: "REPST"
|
||||
seed: 2023
|
||||
|
||||
data:
|
||||
add_day_in_week: true
|
||||
|
|
@ -49,7 +50,6 @@ train:
|
|||
lr_init: 0.003
|
||||
max_grad_norm: 5
|
||||
real_value: true
|
||||
seed: 12
|
||||
weight_decay: 0
|
||||
debug: false
|
||||
output_dim: 1
|
||||
|
|
|
|||
|
|
@ -2,112 +2,88 @@ from utils.normalization import normalize_dataset
|
|||
from dataloader.data_selector import load_st_dataset
|
||||
|
||||
import numpy as np
|
||||
import gc
|
||||
import torch
|
||||
|
||||
|
||||
|
||||
def get_dataloader(args, normalizer="std", single=True):
|
||||
data = load_st_dataset(args) # 加载数据
|
||||
data = load_st_dataset(args)
|
||||
args = args["data"]
|
||||
L, N, F = data.shape # 数据形状
|
||||
L, N, F = data.shape
|
||||
|
||||
# Step 1: data -> x,y
|
||||
# Generate sliding windows for main data and add time features
|
||||
x, y = _prepare_data_with_windows(data, args, single)
|
||||
|
||||
# Split data
|
||||
split_fn = split_data_by_days if args["test_ratio"] > 1 else split_data_by_ratio
|
||||
x_train, x_val, x_test = split_fn(x, args["val_ratio"], args["test_ratio"])
|
||||
y_train, y_val, y_test = split_fn(y, args["val_ratio"], args["test_ratio"])
|
||||
|
||||
# Normalize x and y using the same scaler
|
||||
scaler = _normalize_data(x_train, x_val, x_test, args, normalizer)
|
||||
_apply_existing_scaler(y_train, y_val, y_test, scaler, args)
|
||||
|
||||
# Create dataloaders
|
||||
return (
|
||||
_create_dataloader(x_train, y_train, args["batch_size"], True, False),
|
||||
_create_dataloader(x_val, y_val, args["batch_size"], False, False),
|
||||
_create_dataloader(x_test, y_test, args["batch_size"], False, False),
|
||||
scaler
|
||||
)
|
||||
|
||||
|
||||
def _prepare_data_with_windows(data, args, single):
|
||||
# Generate sliding windows for main data
|
||||
x = add_window_x(data, args["lag"], args["horizon"], single)
|
||||
y = add_window_y(data, args["lag"], args["horizon"], single)
|
||||
|
||||
del data
|
||||
gc.collect()
|
||||
# Generate time features
|
||||
time_features = _generate_time_features(data.shape[0], args)
|
||||
|
||||
# Step 2: time_in_day, day_in_week -> day, week
|
||||
# Add time features to x and y
|
||||
x = _add_time_features(x, time_features, args["lag"], args["horizon"], single, add_window_x)
|
||||
y = _add_time_features(y, time_features, args["lag"], args["horizon"], single, add_window_y)
|
||||
|
||||
return x, y
|
||||
|
||||
|
||||
def _generate_time_features(L, args):
|
||||
N = args["num_nodes"]
|
||||
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 = [(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)
|
||||
return time_in_day, day_in_week
|
||||
|
||||
# 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()
|
||||
def _add_time_features(data, time_features, lag, horizon, single, window_fn):
|
||||
time_in_day, day_in_week = time_features
|
||||
time_day = window_fn(time_in_day, lag, horizon, single)
|
||||
time_week = window_fn(day_in_week, lag, horizon, single)
|
||||
return np.concatenate([data, time_day, time_week], axis=-1)
|
||||
|
||||
# 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()
|
||||
def _normalize_data(train_data, val_data, test_data, args, normalizer):
|
||||
scaler = normalize_dataset(train_data[..., : args["input_dim"]], normalizer, args["column_wise"])
|
||||
|
||||
# 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"]]
|
||||
)
|
||||
for data in [train_data, val_data, test_data]:
|
||||
data[..., : args["input_dim"]] = scaler.transform(data[..., : 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)
|
||||
return scaler
|
||||
|
||||
del time_in_day, day_in_week
|
||||
gc.collect()
|
||||
|
||||
y = np.concatenate([y, y_day, y_week], axis=-1)
|
||||
def _apply_existing_scaler(train_data, val_data, test_data, scaler, args):
|
||||
for data in [train_data, val_data, test_data]:
|
||||
data[..., : args["input_dim"]] = scaler.transform(data[..., : args["input_dim"]])
|
||||
|
||||
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, y_train, args["batch_size"], shuffle=True, drop_last=True
|
||||
)
|
||||
|
||||
del x_train, y_train
|
||||
gc.collect()
|
||||
|
||||
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, 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 _create_dataloader(X_data, Y_data, batch_size, shuffle, drop_last):
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
X_tensor = torch.tensor(X_data, dtype=torch.float32, device=device)
|
||||
Y_tensor = torch.tensor(Y_data, dtype=torch.float32, device=device)
|
||||
dataset = torch.utils.data.TensorDataset(X_tensor, Y_tensor)
|
||||
return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
|
||||
|
||||
|
||||
def split_data_by_days(data, val_days, test_days, interval=30):
|
||||
|
|
@ -128,17 +104,29 @@ def split_data_by_ratio(data, val_ratio, test_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 _generate_windows(data, window=3, horizon=1, offset=0):
|
||||
"""
|
||||
Internal helper function to generate sliding windows.
|
||||
|
||||
:param data: Input data
|
||||
:param window: Window size
|
||||
:param horizon: Horizon size
|
||||
:param offset: Offset from window start
|
||||
:return: Windowed data
|
||||
"""
|
||||
length = len(data)
|
||||
end_index = length - horizon - window + 1
|
||||
windows = []
|
||||
index = 0
|
||||
|
||||
while index < end_index:
|
||||
windows.append(data[index + offset : index + offset + window])
|
||||
index += 1
|
||||
|
||||
return np.array(windows)
|
||||
|
||||
def add_window_x(data, window=3, horizon=1, single=False):
|
||||
"""
|
||||
Generate windowed X values from the input data.
|
||||
|
|
@ -149,17 +137,7 @@ def add_window_x(data, window=3, horizon=1, single=False):
|
|||
: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)
|
||||
|
||||
return _generate_windows(data, window, horizon, offset=0)
|
||||
|
||||
def add_window_y(data, window=3, horizon=1, single=False):
|
||||
"""
|
||||
|
|
@ -171,21 +149,10 @@ def add_window_y(data, window=3, horizon=1, single=False):
|
|||
: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)
|
||||
|
||||
offset = window if not single else window + horizon - 1
|
||||
return _generate_windows(data, window=1 if single else horizon, horizon=horizon, offset=offset)
|
||||
|
||||
if __name__ == "__main__":
|
||||
res = load_st_dataset("SD", 1)
|
||||
k = 1
|
||||
from dataloader.data_selector import load_st_dataset
|
||||
res = load_st_dataset({"dataset": "SD"})
|
||||
print(f"Dataset shape: {res.shape}")
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
from tkinter import Y
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2Model
|
||||
|
|
|
|||
|
|
@ -30,7 +30,6 @@ class TokenEmbedding(nn.Module):
|
|||
|
||||
def forward(self, x):
|
||||
b, n, m, pn, pl = x.shape # batch, node, feature, patch_num, patch_len
|
||||
# 为什么没permute后reshape?
|
||||
x = x.permute(0, 1, 4, 3, 2)
|
||||
x = self.tokenConv(x.reshape(b*n, pl, m*pn)) # batch*node, patch_len, feature*patch_num
|
||||
x = self.confusion_layer(x)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
from tkinter import Y
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2Model
|
||||
|
|
|
|||
2
run.py
2
run.py
|
|
@ -13,7 +13,7 @@ from trainer.trainer_selector import select_trainer
|
|||
def main():
|
||||
args = parse_args()
|
||||
args = init.init_device(args)
|
||||
init.init_seed(args["train"]["seed"])
|
||||
init.init_seed(args["basic"]["seed"])
|
||||
model = init.init_model(args)
|
||||
|
||||
# Load dataset
|
||||
|
|
|
|||
|
|
@ -10,11 +10,14 @@ from tqdm import tqdm
|
|||
|
||||
|
||||
class TrainingStats:
|
||||
"""记录训练过程中的统计信息"""
|
||||
|
||||
def __init__(self, device):
|
||||
self.device = device
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
"""重置所有统计数据"""
|
||||
self.gpu_mem_usage_list = []
|
||||
self.cpu_mem_usage_list = []
|
||||
self.train_time_list = []
|
||||
|
|
@ -24,9 +27,11 @@ class TrainingStats:
|
|||
self.end_time = None
|
||||
|
||||
def start_training(self):
|
||||
"""记录训练开始时间"""
|
||||
self.start_time = time.time()
|
||||
|
||||
def end_training(self):
|
||||
"""记录训练结束时间"""
|
||||
self.end_time = time.time()
|
||||
|
||||
def record_step_time(self, duration, mode):
|
||||
|
|
@ -51,6 +56,10 @@ class TrainingStats:
|
|||
self.cpu_mem_usage_list.append(cpu_mem)
|
||||
self.gpu_mem_usage_list.append(gpu_mem)
|
||||
|
||||
def _calculate_average(self, values_list):
|
||||
"""安全计算平均值，避免除零错误"""
|
||||
return sum(values_list) / len(values_list) if values_list else 0
|
||||
|
||||
def report(self, logger):
|
||||
"""在训练结束时输出汇总统计"""
|
||||
if not self.start_time or not self.end_time:
|
||||
|
|
@ -58,26 +67,10 @@ class TrainingStats:
|
|||
return
|
||||
|
||||
total_time = self.end_time - self.start_time
|
||||
avg_gpu_mem = (
|
||||
sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list)
|
||||
if self.gpu_mem_usage_list
|
||||
else 0
|
||||
)
|
||||
avg_cpu_mem = (
|
||||
sum(self.cpu_mem_usage_list) / len(self.cpu_mem_usage_list)
|
||||
if self.cpu_mem_usage_list
|
||||
else 0
|
||||
)
|
||||
avg_train_time = (
|
||||
sum(self.train_time_list) / len(self.train_time_list)
|
||||
if self.train_time_list
|
||||
else 0
|
||||
)
|
||||
avg_infer_time = (
|
||||
sum(self.infer_time_list) / len(self.infer_time_list)
|
||||
if self.infer_time_list
|
||||
else 0
|
||||
)
|
||||
avg_gpu_mem = self._calculate_average(self.gpu_mem_usage_list)
|
||||
avg_cpu_mem = self._calculate_average(self.cpu_mem_usage_list)
|
||||
avg_train_time = self._calculate_average(self.train_time_list)
|
||||
avg_infer_time = self._calculate_average(self.infer_time_list)
|
||||
iters_per_sec = self.total_iters / total_time if total_time > 0 else 0
|
||||
|
||||
logger.info("===== Training Summary =====")
|
||||
|
|
@ -93,6 +86,8 @@ class TrainingStats:
|
|||
|
||||
|
||||
class Trainer:
|
||||
"""模型训练器，负责整个训练流程的管理"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
|
|
@ -105,26 +100,42 @@ class Trainer:
|
|||
args,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
args = args["train"]
|
||||
train_args = args["train"]
|
||||
|
||||
# 模型和训练相关组件
|
||||
self.model = model
|
||||
self.loss = loss
|
||||
self.optimizer = optimizer
|
||||
self.lr_scheduler = lr_scheduler
|
||||
|
||||
# 数据加载器
|
||||
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.args = train_args
|
||||
|
||||
# 统计信息
|
||||
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._initialize_paths(train_args)
|
||||
self._initialize_logger(train_args)
|
||||
self._initialize_stats()
|
||||
|
||||
def _initialize_paths(self, args):
|
||||
"""初始化模型保存路径"""
|
||||
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
|
||||
def _initialize_logger(self, args):
|
||||
"""初始化日志记录器"""
|
||||
if not os.path.isdir(args["log_dir"]) and not args["debug"]:
|
||||
os.makedirs(args["log_dir"], exist_ok=True)
|
||||
self.logger = get_logger(
|
||||
|
|
@ -132,10 +143,13 @@ class Trainer:
|
|||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
|
||||
# Stats tracker
|
||||
def _initialize_stats(self):
|
||||
"""初始化统计信息记录器"""
|
||||
self.stats = TrainingStats(device=self.device)
|
||||
|
||||
def _run_epoch(self, epoch, dataloader, mode):
|
||||
"""运行一个训练/验证/测试epoch"""
|
||||
# 设置模型模式和是否进行优化
|
||||
if mode == "train":
|
||||
self.model.train()
|
||||
optimizer_step = True
|
||||
|
|
@ -143,48 +157,64 @@ class Trainer:
|
|||
self.model.eval()
|
||||
optimizer_step = False
|
||||
|
||||
# 初始化变量
|
||||
total_loss = 0
|
||||
epoch_time = time.time()
|
||||
|
||||
y_pred, y_true = [], []
|
||||
|
||||
# 训练/验证循环
|
||||
with torch.set_grad_enabled(optimizer_step):
|
||||
progress_bar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}")
|
||||
for batch_idx, (data, target) in progress_bar:
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
|
||||
for _, (data, target) in progress_bar:
|
||||
# 记录步骤开始时间
|
||||
start_time = time.time()
|
||||
|
||||
# 前向传播
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(data).to(self.device)
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
|
||||
loss = self.loss(output, label)
|
||||
|
||||
# 反归一化
|
||||
self.scaler.inverse_transform(output)
|
||||
self.scaler.inverse_transform(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()
|
||||
|
||||
# 记录步骤时间和内存使用
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
|
||||
# 累积损失和预测结果
|
||||
total_loss += loss.item()
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
# Update progress bar with the current loss
|
||||
# 更新进度条
|
||||
progress_bar.set_postfix(loss=loss.item())
|
||||
|
||||
# 合并所有批次的预测结果
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
|
||||
# 计算平均损失
|
||||
avg_loss = total_loss / len(dataloader)
|
||||
# 输出指标
|
||||
|
||||
# 计算并记录指标
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
|
||||
)
|
||||
|
|
@ -192,7 +222,7 @@ class Trainer:
|
|||
f"Epoch #{epoch:02d}: {mode.capitalize():<5} MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s"
|
||||
)
|
||||
|
||||
# 记录内存
|
||||
# 记录内存使用情况
|
||||
self.stats.record_memory_usage()
|
||||
|
||||
return avg_loss
|
||||
|
|
@ -207,22 +237,29 @@ class Trainer:
|
|||
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.stats.start_training()
|
||||
self.logger.info("Training process started")
|
||||
|
||||
# 训练循环
|
||||
for epoch in range(1, self.args["epochs"] + 1):
|
||||
# 训练、验证和测试一个epoch
|
||||
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
|
||||
|
|
@ -231,32 +268,51 @@ class Trainer:
|
|||
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."
|
||||
)
|
||||
# 检查早停条件
|
||||
if self._should_early_stop(not_improved_count):
|
||||
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._save_best_models(best_model, best_test_model)
|
||||
|
||||
# 结束训练并输出统计信息
|
||||
self.stats.end_training()
|
||||
self.stats.report(self.logger)
|
||||
|
||||
# 最终评估
|
||||
self._finalize_training(best_model, best_test_model)
|
||||
|
||||
# 输出参数量
|
||||
# 输出模型参数量
|
||||
self._log_model_params()
|
||||
|
||||
def _should_early_stop(self, not_improved_count):
|
||||
"""检查是否满足早停条件"""
|
||||
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."
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _save_best_models(self, best_model, best_test_model):
|
||||
"""保存最佳模型到文件"""
|
||||
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}"
|
||||
)
|
||||
|
||||
def _log_model_params(self):
|
||||
"""输出模型可训练参数数量"""
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
|
|
@ -276,14 +332,20 @@ class Trainer:
|
|||
|
||||
@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["basic"]["device"])
|
||||
|
||||
# 设置为评估模式
|
||||
model.eval()
|
||||
|
||||
# 收集预测和真实标签
|
||||
y_pred, y_true = [], []
|
||||
|
||||
# 不计算梯度的情况下进行预测
|
||||
with torch.no_grad():
|
||||
for data, target in data_loader:
|
||||
label = target[..., : args["output_dim"]]
|
||||
|
|
@ -291,12 +353,14 @@ class Trainer:
|
|||
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)
|
||||
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
|
|
@ -308,6 +372,7 @@ class Trainer:
|
|||
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"]
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue