trainer修改
This commit is contained in:
parent
f64144f5c1
commit
d4ee8e309e
|
|
@ -0,0 +1,4 @@
|
|||
[mypy]
|
||||
explicit_package_bases = True
|
||||
ignore_missing_imports = True
|
||||
no_site_packages = True
|
||||
|
|
@ -2,6 +2,7 @@ import math
|
|||
import os
|
||||
import time
|
||||
import copy
|
||||
import psutil
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
|
|
@ -23,34 +24,56 @@ class Trainer:
|
|||
args,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
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(
|
||||
args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
# Stats tracker
|
||||
self.stats = TrainingStats(device=args["device"])
|
||||
|
||||
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
|
||||
|
|
@ -58,54 +81,77 @@ 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):
|
||||
with tqdm(
|
||||
total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
) as pbar:
|
||||
for batch_idx, (data, target) in enumerate(dataloader):
|
||||
start_time = time.time()
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(data, labels=label.clone()).to(
|
||||
self.args["device"]
|
||||
)
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
|
||||
for _, (data, target) in progress_bar:
|
||||
# 记录步骤开始时间
|
||||
start_time = time.time()
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
label = self.scaler.inverse_transform(label)
|
||||
# 前向传播
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(data, labels=label.clone()).to(self.device)
|
||||
loss = self.loss(output, label)
|
||||
|
||||
loss = self.loss(output, label)
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
# 反向传播和优化（仅在训练模式）
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += loss.item()
|
||||
|
||||
if mode == "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}"
|
||||
# 梯度裁剪（如果需要）
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
pbar.update(1)
|
||||
pbar.set_postfix(loss=loss.item())
|
||||
# 记录步骤时间和内存使用
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
|
||||
# 累积损失和预测结果
|
||||
total_loss += d_loss.item()
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
# 更新进度条
|
||||
progress_bar.set_postfix(loss=d_loss.item())
|
||||
|
||||
# 合并所有批次的预测结果
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
|
||||
# 计算平均损失
|
||||
avg_loss = total_loss / len(dataloader)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
|
||||
# 计算并记录指标
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
|
||||
)
|
||||
# 记录内存
|
||||
self.logger.info(
|
||||
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
|
||||
|
||||
def train_epoch(self, epoch):
|
||||
|
|
@ -118,21 +164,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
|
||||
|
|
@ -141,38 +195,55 @@ 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)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终评估
|
||||
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):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
self.logger.info("Testing on best validation model")
|
||||
|
|
@ -184,44 +255,44 @@ 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["device"])
|
||||
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"]]
|
||||
output = model(data, labels=label.clone()).to(args["device"])
|
||||
y_pred.append(output)
|
||||
y_true.append(label)
|
||||
output = model(data, labels=label.clone())
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
if args["real_value"]:
|
||||
y_pred = scaler.inverse_transform(y_pred)
|
||||
y_true = scaler.inverse_transform(y_true)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
|
||||
|
||||
for t in range(y_true.shape[1]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -23,44 +23,65 @@ class Trainer:
|
|||
global_config,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = global_config["basic"]["device"]
|
||||
train_config = global_config["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 = train_config
|
||||
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(train_config["log_dir"], "best_model.pth")
|
||||
self.best_test_path = os.path.join(
|
||||
train_config["log_dir"], "best_test_model.pth"
|
||||
)
|
||||
self.loss_figure_path = os.path.join(train_config["log_dir"], "loss.png")
|
||||
|
||||
# Initialize logger
|
||||
if not os.path.isdir(train_config["log_dir"]) and not train_config["debug"]:
|
||||
os.makedirs(train_config["log_dir"], exist_ok=True)
|
||||
# 初始化路径、日志和统计
|
||||
self._initialize_paths(train_config)
|
||||
self._initialize_logger(train_config)
|
||||
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")
|
||||
|
||||
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(
|
||||
train_config["log_dir"],
|
||||
args["log_dir"],
|
||||
name=self.model.__class__.__name__,
|
||||
debug=train_config["debug"],
|
||||
debug=args["debug"],
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {train_config['log_dir']}")
|
||||
# Stats tracker
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
|
||||
def _initialize_stats(self):
|
||||
"""初始化统计信息记录器"""
|
||||
self.stats = TrainingStats(device=self.device)
|
||||
|
||||
def _run_epoch(self, epoch, dataloader, mode):
|
||||
"""运行一个训练/验证/测试epoch"""
|
||||
# 设置模型模式和是否进行优化
|
||||
is_train = mode == "train"
|
||||
self.model.train() if is_train else self.model.eval()
|
||||
|
||||
# 初始化变量
|
||||
total_loss = 0.0
|
||||
epoch_time = time.time()
|
||||
y_pred, y_true = [], []
|
||||
|
||||
with (
|
||||
torch.set_grad_enabled(is_train),
|
||||
|
|
@ -85,10 +106,12 @@ class Trainer:
|
|||
|
||||
# compute loss
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
loss = self.loss(output, label)
|
||||
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
# backward / step
|
||||
if is_train:
|
||||
loss.backward()
|
||||
|
|
@ -98,22 +121,39 @@ class Trainer:
|
|||
)
|
||||
self.optimizer.step()
|
||||
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += loss.item()
|
||||
total_loss += d_loss.item()
|
||||
|
||||
# 累积预测结果
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
# 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}"
|
||||
f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
|
||||
)
|
||||
|
||||
pbar.update(1)
|
||||
pbar.set_postfix(loss=loss.item())
|
||||
pbar.set_postfix(loss=d_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"]
|
||||
)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
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()
|
||||
|
|
@ -129,21 +169,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
|
||||
|
|
@ -152,38 +200,55 @@ 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)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终评估
|
||||
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):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
self.logger.info("Testing on best validation model")
|
||||
|
|
@ -195,51 +260,44 @@ class Trainer:
|
|||
|
||||
@staticmethod
|
||||
def test(model, args, data_loader, scaler, logger, path=None):
|
||||
global_config = args
|
||||
device = global_config["basic"]["device"]
|
||||
args = global_config["train"]
|
||||
"""对模型进行评估并输出性能指标"""
|
||||
# 加载模型检查点（如果提供了路径）
|
||||
if path:
|
||||
checkpoint = torch.load(path)
|
||||
model.load_state_dict(checkpoint["state_dict"])
|
||||
model.to(device)
|
||||
model.to(args["basic"]["device"])
|
||||
|
||||
# 设置为评估模式
|
||||
model.eval()
|
||||
|
||||
# 收集预测和真实标签
|
||||
y_pred, y_true = [], []
|
||||
|
||||
# 不计算梯度的情况下进行预测
|
||||
with torch.no_grad():
|
||||
for data, target, cycle_index in data_loader:
|
||||
label = target[..., : args["output_dim"]]
|
||||
output = model(data, cycle_index)
|
||||
y_pred.append(output)
|
||||
y_true.append(label)
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
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)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(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]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import math
|
|||
import os
|
||||
import time
|
||||
import copy
|
||||
import psutil
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
|
|
@ -23,34 +24,56 @@ class Trainer:
|
|||
args,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
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(
|
||||
args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
# Stats tracker
|
||||
self.stats = TrainingStats(device=args["device"])
|
||||
|
||||
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
|
||||
|
|
@ -58,52 +81,77 @@ 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):
|
||||
with tqdm(
|
||||
total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
) as pbar:
|
||||
for batch_idx, (data, target) in enumerate(dataloader):
|
||||
start_time = time.time()
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(data).to(self.args["device"])
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
|
||||
for _, (data, target) in progress_bar:
|
||||
# 记录步骤开始时间
|
||||
start_time = time.time()
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
# 前向传播
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(data).to(self.device)
|
||||
loss = self.loss(output, label)
|
||||
|
||||
loss = self.loss(output, label)
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
# 反向传播和优化（仅在训练模式）
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
|
||||
total_loss += loss.item()
|
||||
|
||||
if mode == "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}"
|
||||
# 梯度裁剪（如果需要）
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
pbar.update(1)
|
||||
pbar.set_postfix(loss=loss.item())
|
||||
# 记录步骤时间和内存使用
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
|
||||
# 累积损失和预测结果
|
||||
total_loss += d_loss.item()
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
# 更新进度条
|
||||
progress_bar.set_postfix(loss=d_loss.item())
|
||||
|
||||
# 合并所有批次的预测结果
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
|
||||
# 计算平均损失
|
||||
avg_loss = total_loss / len(dataloader)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
|
||||
# 计算并记录指标
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
|
||||
)
|
||||
# 记录内存
|
||||
self.logger.info(
|
||||
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
|
||||
|
||||
def train_epoch(self, epoch):
|
||||
|
|
@ -116,21 +164,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
|
||||
|
|
@ -139,37 +195,55 @@ 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)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终评估
|
||||
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):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
self.logger.info("Testing on best validation model")
|
||||
|
|
@ -181,48 +255,44 @@ 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["device"])
|
||||
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"]]
|
||||
output = model(data)
|
||||
y_pred.append(output)
|
||||
y_true.append(label)
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
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)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(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]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -23,35 +23,57 @@ class Trainer:
|
|||
args,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
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.batches_seen = 0
|
||||
|
||||
# 统计信息
|
||||
self.train_per_epoch = len(train_loader)
|
||||
self.val_per_epoch = len(val_loader) if val_loader else 0
|
||||
self.batches_seen = 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(
|
||||
args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
# Stats tracker
|
||||
self.stats = TrainingStats(device=args["device"])
|
||||
|
||||
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
|
||||
|
|
@ -59,55 +81,86 @@ 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):
|
||||
with tqdm(
|
||||
total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
) as pbar:
|
||||
for batch_idx, (data, target) in enumerate(dataloader):
|
||||
start_time = time.time()
|
||||
self.batches_seen += 1
|
||||
label = target[..., : self.args["output_dim"]].clone()
|
||||
output = self.model(data, target, self.batches_seen).to(
|
||||
self.args["device"]
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
|
||||
for batch_idx, (data, target) in progress_bar:
|
||||
start_time = time.time()
|
||||
self.batches_seen += 1
|
||||
label = target[..., : self.args["output_dim"]].clone()
|
||||
|
||||
# 前向传播
|
||||
if mode == "train":
|
||||
output = self.model(data, target, self.batches_seen).to(self.device)
|
||||
else:
|
||||
output = self.model(data, target).to(self.device)
|
||||
|
||||
# 计算原始loss
|
||||
loss = self.loss(output, label)
|
||||
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_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 += d_loss.item()
|
||||
|
||||
# 累积预测结果
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0:
|
||||
self.logger.info(
|
||||
f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
|
||||
)
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
# 更新 tqdm 的进度
|
||||
progress_bar.update(1)
|
||||
progress_bar.set_postfix(loss=d_loss.item())
|
||||
|
||||
loss = self.loss(output, 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()
|
||||
|
||||
# record step time
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += loss.item()
|
||||
|
||||
if mode == "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}"
|
||||
)
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
pbar.update(1)
|
||||
pbar.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)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
|
||||
# 计算并记录指标
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
|
||||
)
|
||||
# 记录内存
|
||||
self.logger.info(
|
||||
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
|
||||
|
||||
def train_epoch(self, epoch):
|
||||
|
|
@ -120,21 +173,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
|
||||
|
|
@ -143,37 +204,54 @@ 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)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 输出模型参数量
|
||||
self._log_model_params()
|
||||
|
||||
# 最终评估
|
||||
self._finalize_training(best_model, best_test_model)
|
||||
|
||||
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):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
|
|
@ -186,44 +264,44 @@ 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["device"])
|
||||
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"]].clone()
|
||||
output = model(data, target)
|
||||
y_pred.append(output)
|
||||
y_true.append(label)
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
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)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
|
||||
|
||||
for t in range(y_true.shape[1]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -26,42 +26,35 @@ class Trainer:
|
|||
args,
|
||||
lr_scheduler=None,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
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["train"]
|
||||
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.best_path = os.path.join(self.args["log_dir"], "best_model.pth")
|
||||
self.best_test_path = os.path.join(self.args["log_dir"], "best_test_model.pth")
|
||||
self.loss_figure_path = os.path.join(self.args["log_dir"], "loss.png")
|
||||
self.pretrain_dir = (
|
||||
f"./pre-train/{args['model']['type']}/{args['data']['type']}"
|
||||
)
|
||||
self.pretrain_path = os.path.join(self.pretrain_dir, "best_model.pth")
|
||||
self.pretrain_best_path = os.path.join(self.pretrain_dir, "best_test_model.pth")
|
||||
|
||||
# Initialize logger
|
||||
if not os.path.isdir(self.args["log_dir"]) and not self.args["debug"]:
|
||||
os.makedirs(self.args["log_dir"], exist_ok=True)
|
||||
if not os.path.isdir(self.pretrain_dir) and not self.args["debug"]:
|
||||
os.makedirs(self.pretrain_dir, exist_ok=True)
|
||||
self.logger = get_logger(
|
||||
self.args["log_dir"],
|
||||
name=self.model.__class__.__name__,
|
||||
debug=self.args["debug"],
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {self.args['log_dir']}")
|
||||
# Stats tracker
|
||||
self.stats = TrainingStats(device=args["device"])
|
||||
|
||||
# 初始化路径、日志和统计
|
||||
self._initialize_paths(args, train_args)
|
||||
self._initialize_logger(train_args)
|
||||
self._initialize_stats()
|
||||
|
||||
# 教师-学生蒸馏相关
|
||||
if self.args["teacher_stu"]:
|
||||
self.tmodel = self.loadTeacher(args)
|
||||
else:
|
||||
|
|
@ -70,9 +63,41 @@ class Trainer:
|
|||
f"./pre-train/{args['model']['type']}/{args['data']['type']}/best_model.pth"
|
||||
f"然后在config中配置train.teacher_stu模式为True开启蒸馏模式"
|
||||
)
|
||||
|
||||
def _initialize_paths(self, args, train_args):
|
||||
"""初始化模型保存路径"""
|
||||
self.best_path = os.path.join(train_args["log_dir"], "best_model.pth")
|
||||
self.best_test_path = os.path.join(train_args["log_dir"], "best_test_model.pth")
|
||||
self.loss_figure_path = os.path.join(train_args["log_dir"], "loss.png")
|
||||
self.pretrain_dir = (
|
||||
f"./pre-train/{args['model']['type']}/{args['data']['type']}"
|
||||
)
|
||||
self.pretrain_path = os.path.join(self.pretrain_dir, "best_model.pth")
|
||||
self.pretrain_best_path = os.path.join(self.pretrain_dir, "best_test_model.pth")
|
||||
|
||||
# 创建预训练目录
|
||||
if not os.path.isdir(self.pretrain_dir) and not train_args["debug"]:
|
||||
os.makedirs(self.pretrain_dir, exist_ok=True)
|
||||
|
||||
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(
|
||||
args["log_dir"],
|
||||
name=self.model.__class__.__name__,
|
||||
debug=args["debug"],
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
|
||||
def _initialize_stats(self):
|
||||
"""初始化统计信息记录器"""
|
||||
self.stats = TrainingStats(device=self.device)
|
||||
|
||||
def _run_epoch(self, epoch, dataloader, mode):
|
||||
"""运行一个训练/验证/测试epoch"""
|
||||
# self.tmodel.eval()
|
||||
# 设置模型模式和是否进行优化
|
||||
if mode == "train":
|
||||
self.model.train()
|
||||
optimizer_step = True
|
||||
|
|
@ -80,8 +105,10 @@ 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):
|
||||
with tqdm(
|
||||
|
|
@ -89,15 +116,17 @@ class Trainer:
|
|||
) as pbar:
|
||||
for batch_idx, (data, target) in enumerate(dataloader):
|
||||
start_time = time.time()
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
|
||||
if self.args["teacher_stu"]:
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
# 教师-学生蒸馏模式
|
||||
output, out_, _ = self.model(data)
|
||||
gout, tout, sout = self.tmodel(data)
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
|
||||
|
||||
# 计算原始loss
|
||||
loss1 = self.loss(output, label)
|
||||
|
||||
# 计算蒸馏相关loss
|
||||
scl = self.loss_cls(out_, sout)
|
||||
kl_loss = nn.KLDivLoss(
|
||||
reduction="batchmean", log_target=True
|
||||
|
|
@ -105,17 +134,22 @@ class Trainer:
|
|||
gout = F.log_softmax(gout, dim=-1).cuda()
|
||||
mlp_emb_ = F.log_softmax(output, dim=-1).cuda()
|
||||
tkloss = kl_loss(mlp_emb_.cuda().float(), gout.cuda().float())
|
||||
|
||||
# 总loss
|
||||
loss = loss1 + 10 * tkloss + 1 * scl
|
||||
|
||||
else:
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
# 普通训练模式
|
||||
output, out_, _ = self.model(data)
|
||||
|
||||
if self.args["real_value"]:
|
||||
output = self.scaler.inverse_transform(output)
|
||||
|
||||
loss = self.loss(output, label)
|
||||
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
# 反向传播和优化（仅在训练模式）
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
|
@ -128,20 +162,34 @@ class Trainer:
|
|||
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += loss.item()
|
||||
total_loss += d_loss.item()
|
||||
|
||||
# 累积预测结果
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
if mode == "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}"
|
||||
f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
|
||||
)
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
pbar.update(1)
|
||||
pbar.set_postfix(loss=loss.item())
|
||||
pbar.set_postfix(loss=d_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"]
|
||||
)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
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()
|
||||
|
|
@ -157,6 +205,7 @@ 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
|
||||
|
|
@ -182,13 +231,7 @@ 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:
|
||||
|
|
@ -207,14 +250,25 @@ class Trainer:
|
|||
# 输出统计与参数
|
||||
self.stats.end_training()
|
||||
self.stats.report(self.logger)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
self._log_model_params()
|
||||
self._finalize_training(best_model, best_test_model)
|
||||
|
||||
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 _log_model_params(self):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
|
|
@ -274,48 +328,44 @@ 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["device"])
|
||||
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"]]
|
||||
output, _, _ = model(data)
|
||||
y_pred.append(output)
|
||||
y_true.append(label)
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
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)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(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/PEMSD8/y_true.pt") # [3566,12,170,1]
|
||||
|
||||
for t in range(y_true.shape[1]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -25,37 +25,60 @@ class Trainer:
|
|||
times,
|
||||
w,
|
||||
):
|
||||
# 设备和基本参数
|
||||
self.device = args["basic"]["device"]
|
||||
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
|
||||
self.device = args["device"]
|
||||
|
||||
# Paths for saving models and logs
|
||||
|
||||
# 初始化路径、日志和统计
|
||||
self._initialize_paths(train_args)
|
||||
self._initialize_logger(train_args)
|
||||
self._initialize_stats()
|
||||
|
||||
# 模型特定参数
|
||||
self.times = times.to(self.device, dtype=torch.float)
|
||||
self.w = w
|
||||
|
||||
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(
|
||||
args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
|
||||
)
|
||||
self.logger.info(f"Experiment log path in: {args['log_dir']}")
|
||||
# Stats tracker
|
||||
self.stats = TrainingStats(device=args["device"])
|
||||
self.times = times.to(self.device, dtype=torch.float)
|
||||
self.w = w
|
||||
|
||||
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
|
||||
|
|
@ -63,53 +86,84 @@ 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):
|
||||
with tqdm(
|
||||
total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
) as pbar:
|
||||
for batch_idx, batch in enumerate(dataloader):
|
||||
start_time = time.time()
|
||||
batch = tuple(b.to(self.device, dtype=torch.float) for b in batch)
|
||||
*train_coeffs, target = batch
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
output = self.model(self.times, train_coeffs)
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
|
||||
for batch_idx, batch in progress_bar:
|
||||
start_time = time.time()
|
||||
batch = tuple(b.to(self.device, dtype=torch.float) for b in batch)
|
||||
*train_coeffs, target = batch
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
|
||||
# 前向传播
|
||||
output = self.model(self.times, train_coeffs)
|
||||
|
||||
# 计算原始loss
|
||||
loss = self.loss(output, label)
|
||||
|
||||
# if self.args['real_value']:
|
||||
# output = self.scaler.inverse_transform(output)
|
||||
# 反归一化
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
|
||||
loss = self.loss(output, label)
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
# 反归一化的loss
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
# 反向传播和优化（仅在训练模式）
|
||||
if optimizer_step and self.optimizer is not None:
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += loss.item()
|
||||
|
||||
if mode == "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}"
|
||||
if self.args["grad_norm"]:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.model.parameters(), self.args["max_grad_norm"]
|
||||
)
|
||||
self.optimizer.step()
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
pbar.update(1)
|
||||
pbar.set_postfix(loss=loss.item())
|
||||
# 记录步骤时间
|
||||
step_time = time.time() - start_time
|
||||
self.stats.record_step_time(step_time, mode)
|
||||
total_loss += d_loss.item()
|
||||
|
||||
# 累积预测结果
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
if mode == "train" and (batch_idx + 1) % self.args["log_step"] == 0:
|
||||
self.logger.info(
|
||||
f"Train Epoch {epoch}: {batch_idx + 1}/{len(dataloader)} Loss: {d_loss.item():.6f}"
|
||||
)
|
||||
|
||||
# 更新 tqdm 的进度
|
||||
progress_bar.update(1)
|
||||
progress_bar.set_postfix(loss=d_loss.item())
|
||||
|
||||
# 合并所有批次的预测结果
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
|
||||
# 计算平均损失
|
||||
avg_loss = total_loss / len(dataloader)
|
||||
self.logger.info(
|
||||
f"{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s"
|
||||
|
||||
# 计算并记录指标
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
|
||||
)
|
||||
# 记录内存
|
||||
self.logger.info(
|
||||
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
|
||||
|
||||
def train_epoch(self, epoch):
|
||||
|
|
@ -122,21 +176,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
|
||||
|
|
@ -145,37 +207,54 @@ 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)
|
||||
try:
|
||||
total_params = sum(
|
||||
p.numel() for p in self.model.parameters() if p.requires_grad
|
||||
)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 输出模型参数量
|
||||
self._log_model_params()
|
||||
|
||||
# 最终评估
|
||||
self._finalize_training(best_model, best_test_model)
|
||||
|
||||
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):
|
||||
"""输出模型可训练参数数量"""
|
||||
total_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad)
|
||||
self.logger.info(f"Trainable params: {total_params}")
|
||||
|
||||
|
||||
def _finalize_training(self, best_model, best_test_model):
|
||||
self.model.load_state_dict(best_model)
|
||||
|
|
@ -188,42 +267,41 @@ class Trainer:
|
|||
|
||||
@staticmethod
|
||||
def test(model, args, data_loader, scaler, logger):
|
||||
"""对模型进行评估并输出性能指标"""
|
||||
# 设置为评估模式
|
||||
model.eval()
|
||||
|
||||
# 收集预测和真实标签
|
||||
y_pred, y_true = [], []
|
||||
times = torch.linspace(0, 11, 12)
|
||||
|
||||
# 不计算梯度的情况下进行预测
|
||||
with torch.no_grad():
|
||||
for batch_idx, batch in enumerate(data_loader):
|
||||
batch = tuple(b.to(args["device"], dtype=torch.float) for b in batch)
|
||||
batch = tuple(b.to(args["basic"]["device"], dtype=torch.float) for b in batch)
|
||||
*test_coeffs, target = batch
|
||||
label = target[..., : args["output_dim"]]
|
||||
output = model(times.to(args["device"], dtype=torch.float), test_coeffs)
|
||||
y_true.append(label)
|
||||
y_pred.append(output)
|
||||
output = model(times.to(args["basic"]["device"], dtype=torch.float), test_coeffs)
|
||||
y_true.append(label.detach().cpu())
|
||||
y_pred.append(output.detach().cpu())
|
||||
|
||||
# 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)
|
||||
# 反归一化
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
|
||||
d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
|
||||
|
||||
for t in range(y_true.shape[1]):
|
||||
# 计算并记录每个时间步的指标
|
||||
for t in range(d_y_true.shape[1]):
|
||||
mae, rmse, mape = all_metrics(
|
||||
y_pred[:, t, ...],
|
||||
y_true[:, t, ...],
|
||||
d_y_pred[:, t, ...],
|
||||
d_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}"
|
||||
)
|
||||
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}"
|
||||
)
|
||||
# 计算并记录平均指标
|
||||
mae, rmse, mape = all_metrics(d_y_pred, d_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):
|
||||
|
|
|
|||
|
|
@ -47,6 +47,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:
|
||||
|
|
@ -54,26 +58,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 =====")
|
||||
|
|
|
|||
Loading…
Reference in New Issue