czzhangheng
/
TrafficWheel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

opt trainer

This commit is contained in:
czzhangheng 2025-12-15 00:38:43 +08:00
parent 3b4acd4951
commit 97743dfd05
2 changed files with 78 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
train.py
View File

@ -6,14 +6,12 @@ import utils.initializer as init
from dataloader.loader_selector import get_dataloader from dataloader.loader_selector import get_dataloader
from trainer.trainer_selector import select_trainer from trainer.trainer_selector import select_trainer
import cProfile
def read_config(config_path): def read_config(config_path):
with open(config_path, "r") as file: with open(config_path, "r") as file:
config = yaml.safe_load(file) config = yaml.safe_load(file)
# 全局配置 # 全局配置
device = "cuda:1" # 指定设备为cuda:0 device = "cpu" # 指定设备为cuda:0
seed = 2023 # 随机种子 seed = 2023 # 随机种子
epochs = 120 epochs = 120
@ -67,8 +65,8 @@ def main(debug=False):
model_list = ["iTransformer"] model_list = ["iTransformer"]
# 指定数据集 # 指定数据集
# dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"] # dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]
# dataset_list = ["AirQuality"] dataset_list = ["AirQuality"]
dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"] # dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"]
# 我的调试开关，不做测试就填 str(False) # 我的调试开关，不做测试就填 str(False)
# os.environ["TRY"] = str(False) # os.environ["TRY"] = str(False)
@ -99,4 +97,4 @@ def main(debug=False):
if __name__ == "__main__": if __name__ == "__main__":
# 调试用 # 调试用
main(debug = False) main(debug = True)

163
trainer/Trainer.py
View File

@ -10,47 +10,45 @@ from tqdm import tqdm
class Trainer: class Trainer:
"""模型训练器，负责整个训练流程的管理""" """模型训练器，负责整个训练流程的管理"""
def __init__(self, model, loss, optimizer, def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args, lr_scheduler=None):
train_loader, val_loader, test_loader, scaler,
args, lr_scheduler=None,):
# 设备和基本参数 # 设备和基本参数
self.config = args self.config = args
self.device = args["basic"]["device"] self.device = args["basic"]["device"]
train_args = args["train"] self.args = args["train"]
# 模型和训练相关组件 # 模型和训练相关组件
self.model = model self.model, self.loss, self.optimizer, self.lr_scheduler = model, loss, optimizer, lr_scheduler
self.loss = loss
self.optimizer = optimizer
self.lr_scheduler = lr_scheduler
# 数据加载器 # 数据加载器
self.train_loader = train_loader self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader
self.val_loader = val_loader
self.test_loader = test_loader
# 数据处理工具 # 数据处理工具
self.scaler = scaler self.scaler = scaler
self.args = train_args
# 初始化路径、日志和统计 # 初始化路径、日志和统计
self._initialize_paths(train_args) self._initialize_paths(self.args)
self._initialize_logger(train_args) self._initialize_logger(self.args)
def _initialize_paths(self, args): def _initialize_paths(self, args):
"""初始化模型保存路径""" """初始化模型保存路径"""
self.best_path = os.path.join(args["log_dir"], "best_model.pth") log_dir = args["log_dir"]
self.best_test_path = os.path.join(args["log_dir"], "best_test_model.pth") self.best_path = os.path.join(log_dir, "best_model.pth")
self.loss_figure_path = os.path.join(args["log_dir"], "loss.png") self.best_test_path = os.path.join(log_dir, "best_test_model.pth")
self.loss_figure_path = os.path.join(log_dir, "loss.png")
def _initialize_logger(self, args): def _initialize_logger(self, args):
"""初始化日志记录器""" """初始化日志记录器"""
if not os.path.isdir(args["log_dir"]) and not args["debug"]: log_dir = args["log_dir"]
os.makedirs(args["log_dir"], exist_ok=True) if not args["debug"]:
self.logger = get_logger(args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]) os.makedirs(log_dir, exist_ok=True)
self.logger.info(f"Experiment log path in: {args['log_dir']}") self.logger = get_logger(log_dir, name=self.model.__class__.__name__, debug=args["debug"])
self.logger.info(f"Experiment log path in: {log_dir}")
def _run_epoch(self, epoch, dataloader, mode): def _run_epoch(self, epoch, dataloader, mode):
"""运行一个训练/验证/测试epoch""" """运行一个训练/验证/测试epoch"""
# 设置模型模式和是否进行优化 # 设置模型模式和是否进行优化
if mode == "train": self.model.train(); optimizer_step = True self.model.train() if mode == "train" else self.model.eval()
else: self.model.eval(); optimizer_step = False optimizer_step = mode == "train"
# 初始化变量 # 初始化变量
total_loss = 0 total_loss = 0
@ -60,105 +58,111 @@ class Trainer:
# 训练/验证循环 # 训练/验证循环
with torch.set_grad_enabled(optimizer_step): with torch.set_grad_enabled(optimizer_step):
progress_bar = tqdm( progress_bar = tqdm(
enumerate(dataloader), dataloader,
total=len(dataloader), total=len(dataloader),
desc=f"{mode.capitalize()} Epoch {epoch}" desc=f"{mode.capitalize()} Epoch {epoch}"
) )
for _, (data, target) in progress_bar: for data, target in progress_bar:
# 转移数据 # 转移数据并提取标签
data = data.to(self.device) data, target = data.to(self.device), target.to(self.device)
target = target.to(self.device)
label = target[..., : self.args["output_dim"]] label = target[..., : self.args["output_dim"]]
# 计算loss和反归一化loss
# 计算输出
output = self.model(data) output = self.model(data)
# 我的调试开关 # 我的调试开关
if os.environ.get("TRY") == "True": if os.environ.get("TRY") == "True":
print(f"[{'✅' if output.shape == label.shape else '❌'}]: output: {output.shape}, label: {label.shape}") status = '✅' if output.shape == label.shape else '❌'
print(f"[{status}]: output: {output.shape}, label: {label.shape}")
assert False assert False
# 计算损失
loss = self.loss(output, label) loss = self.loss(output, label)
d_output = self.scaler.inverse_transform(output) d_output = self.scaler.inverse_transform(output)
d_label = self.scaler.inverse_transform(label) d_label = self.scaler.inverse_transform(label)
d_loss = self.loss(d_output, d_label) d_loss = self.loss(d_output, d_label)
# 累积损失和预测结果 # 累积损失和预测结果
total_loss += d_loss.item() total_loss += d_loss.item()
y_pred.append(d_output.detach().cpu()) y_pred.append(d_output.detach().cpu())
y_true.append(d_label.detach().cpu()) y_true.append(d_label.detach().cpu())
# 反向传播和优化（仅在训练模式） # 反向传播和优化（仅在训练模式）
if optimizer_step and self.optimizer is not None: if optimizer_step and self.optimizer is not None:
self.optimizer.zero_grad() self.optimizer.zero_grad()
loss.backward() loss.backward()
# 梯度裁剪（如果需要） # 梯度裁剪（如果需要）
if self.args["grad_norm"]: if self.args["grad_norm"]:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args["max_grad_norm"]) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args["max_grad_norm"])
self.optimizer.step() self.optimizer.step()
# 更新进度条 # 更新进度条
progress_bar.set_postfix(loss=d_loss.item()) progress_bar.set_postfix(loss=d_loss.item())
# 合并所有批次的预测结果 # 合并所有批次的预测结果
y_pred = torch.cat(y_pred, dim=0) y_pred, y_true = torch.cat(y_pred, dim=0), torch.cat(y_true, dim=0)
y_true = torch.cat(y_true, dim=0)
# 计算损失并记录指标 # 计算损失并记录指标
avg_loss = total_loss / len(dataloader) avg_loss = total_loss / len(dataloader)
mae, rmse, mape = all_metrics(y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]) mae, rmse, mape = all_metrics(y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"])
self.logger.info( self.logger.info(
f"Epoch #{epoch:02d}: {mode.capitalize():<5} " f"Epoch #{epoch:02d}: {mode.capitalize():<5} "
f"MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s" f"MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s"
) )
return avg_loss return avg_loss
def train_epoch(self, epoch):
return self._run_epoch(epoch, self.train_loader, "train")
def val_epoch(self, epoch):
return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
def test_epoch(self, epoch):
return self._run_epoch(epoch, self.test_loader, "test")
def train(self): def train(self):
# 初始化记录 # 初始化记录
best_model, best_test_model = None, None best_model = best_test_model = None
best_loss, best_test_loss = float("inf"), float("inf") best_loss = best_test_loss = float("inf")
not_improved_count = 0 not_improved_count = 0
# 开始训练 # 开始训练
self.logger.info("Training process started") self.logger.info("Training process started")
# 训练循环 # 训练循环
for epoch in range(1, self.args["epochs"] + 1): for epoch in range(1, self.args["epochs"] + 1):
# 训练、验证和测试一个epoch # 训练、验证和测试一个epoch
train_epoch_loss = self.train_epoch(epoch) train_epoch_loss = self._run_epoch(epoch, self.train_loader, "train")
val_epoch_loss = self.val_epoch(epoch) val_epoch_loss = self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
test_epoch_loss = self.test_epoch(epoch) test_epoch_loss = self._run_epoch(epoch, self.test_loader, "test")
# 检查梯度爆炸 # 检查梯度爆炸
if train_epoch_loss > 1e6: if train_epoch_loss > 1e6:
self.logger.warning("Gradient explosion detected. Ending...") self.logger.warning("Gradient explosion detected. Ending...")
break break
# 更新最佳验证模型 # 更新最佳验证模型
if val_epoch_loss < best_loss: if val_epoch_loss < best_loss:
best_loss = val_epoch_loss best_loss, not_improved_count = val_epoch_loss, 0
not_improved_count = 0
best_model = copy.deepcopy(self.model.state_dict()) best_model = copy.deepcopy(self.model.state_dict())
self.logger.info("Best validation model saved!") self.logger.info("Best validation model saved!")
else: else:
not_improved_count += 1 not_improved_count += 1
# 早停
# 早停检查
if self._should_early_stop(not_improved_count): if self._should_early_stop(not_improved_count):
break break
# 更新最佳测试模型 # 更新最佳测试模型
if test_epoch_loss < best_test_loss: if test_epoch_loss < best_test_loss:
best_test_loss = test_epoch_loss best_test_loss = test_epoch_loss
best_test_model = copy.deepcopy(self.model.state_dict()) best_test_model = copy.deepcopy(self.model.state_dict())
# 保存最佳模型 # 保存最佳模型
if not self.args["debug"]: if not self.args["debug"]:
self._save_best_models(best_model, best_test_model) self._save_best_models(best_model, best_test_model)
# 最终评估 # 最终评估
self._finalize_training(best_model, best_test_model) self._finalize_training(best_model, best_test_model)
def _should_early_stop(self, not_improved_count): def _should_early_stop(self, not_improved_count):
"""检查是否满足早停条件""" """检查是否满足早停条件"""
if ( if self.args["early_stop"] and not_improved_count == self.args["early_stop_patience"]:
self.args["early_stop"]
and not_improved_count == self.args["early_stop_patience"]
):
self.logger.info( self.logger.info(
f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops." f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops."
) )
@ -190,59 +194,44 @@ class Trainer:
@staticmethod @staticmethod
def test(model, args, data_loader, scaler, logger, path=None): def test(model, args, data_loader, scaler, logger, path=None):
"""对模型进行评估并输出性能指标""" """对模型进行评估并输出性能指标"""
# 确定设备信息 # 验证参数类型
device = None if not isinstance(args, dict):
output_dim = None
# 处理不同的参数格式
if isinstance(args, dict):
if "basic" in args:
# 完整配置情况
device = args["basic"]["device"]
output_dim = args["train"]["output_dim"]
else:
# 只有train_args情况，从模型获取设备
device = next(model.parameters()).device
output_dim = args["output_dim"]
else:
raise ValueError(f"Unsupported args type: {type(args)}") raise ValueError(f"Unsupported args type: {type(args)}")
# 确定设备和输出维度
is_full_config = "basic" in args
device = args["basic"]["device"] if is_full_config else next(model.parameters()).device
output_dim = args["train"]["output_dim"] if is_full_config else args["output_dim"]
# 获取metrics参数
train_args = args["train"] if is_full_config else args
mae_thresh, mape_thresh = train_args["mae_thresh"], train_args["mape_thresh"]
# 加载模型检查点（如果提供了路径） # 加载模型检查点（如果提供了路径）
if path: if path:
checkpoint = torch.load(path) checkpoint = torch.load(path)
model.load_state_dict(checkpoint["state_dict"]) model.load_state_dict(checkpoint["state_dict"])
model.to(device) model.to(device)
# 设置为评估模式 # 设置为评估模式并收集预测结果
model.eval() model.eval()
# 收集预测和真实标签
y_pred, y_true = [], [] y_pred, y_true = [], []
# 不计算梯度的情况下进行预测 # 不计算梯度的情况下进行预测
with torch.no_grad(): with torch.no_grad():
for data, target in data_loader: for data, target in data_loader:
# 将数据和标签移动到指定设备 # 将数据和标签移动到指定设备
data = data.to(device) data, target = data.to(device), target.to(device)
target = target.to(device)
label = target[..., : output_dim] label = target[..., : output_dim]
output = model(data) output = model(data)
y_pred.append(output.detach().cpu()) y_pred.append(output.detach().cpu())
y_true.append(label.detach().cpu()) y_true.append(label.detach().cpu())
# 反归一化并计算指标
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, 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)) d_y_true = scaler.inverse_transform(torch.cat(y_true, dim=0))
# 获取metrics参数
if "basic" in args:
# 完整配置情况
mae_thresh = args["train"]["mae_thresh"]
mape_thresh = args["train"]["mape_thresh"]
else:
# 只有train_args情况
mae_thresh = args["mae_thresh"]
mape_thresh = args["mape_thresh"]
# 计算并记录每个时间步的指标 # 计算并记录每个时间步的指标
for t in range(d_y_true.shape[1]): for t in range(d_y_true.shape[1]):
mae, rmse, mape = all_metrics( mae, rmse, mape = all_metrics(
@ -254,9 +243,5 @@ class Trainer:
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(d_y_pred, d_y_true, mae_thresh, mape_thresh) avg_mae, avg_rmse, avg_mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh)
logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}") logger.info(f"Average Horizon, MAE: {avg_mae:.4f}, RMSE: {avg_rmse:.4f}, MAPE: {avg_mape:.4f}")
@staticmethod
def _compute_sampling_threshold(global_step, k):
return k / (k + math.exp(global_step / k))