opt trainer
This commit is contained in:
parent
3b4acd4951
commit
97743dfd05
10
train.py
10
train.py
|
|
@ -6,14 +6,12 @@ import utils.initializer as init
|
|||
from dataloader.loader_selector import get_dataloader
|
||||
from trainer.trainer_selector import select_trainer
|
||||
|
||||
import cProfile
|
||||
|
||||
def read_config(config_path):
|
||||
with open(config_path, "r") as file:
|
||||
config = yaml.safe_load(file)
|
||||
|
||||
# 全局配置
|
||||
device = "cuda:1" # 指定设备为cuda:0
|
||||
device = "cpu" # 指定设备为cuda:0
|
||||
seed = 2023 # 随机种子
|
||||
epochs = 120
|
||||
|
||||
|
|
@ -67,8 +65,8 @@ def main(debug=False):
|
|||
model_list = ["iTransformer"]
|
||||
# 指定数据集
|
||||
# dataset_list = ["AirQuality", "SolarEnergy", "PEMS-BAY", "METR-LA", "BJTaxi-Inflow", "BJTaxi-Outflow", "NYCBike-Inflow", "NYCBike-Outflow"]
|
||||
# dataset_list = ["AirQuality"]
|
||||
dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"]
|
||||
dataset_list = ["AirQuality"]
|
||||
# dataset_list = ["AirQuality", "SolarEnergy", "METR-LA", "NYCBike-Inflow", "NYCBike-Outflow"]
|
||||
|
||||
# 我的调试开关，不做测试就填 str(False)
|
||||
# os.environ["TRY"] = str(False)
|
||||
|
|
@ -99,4 +97,4 @@ def main(debug=False):
|
|||
|
||||
if __name__ == "__main__":
|
||||
# 调试用
|
||||
main(debug = False)
|
||||
main(debug = True)
|
||||
|
|
@ -10,47 +10,45 @@ from tqdm import tqdm
|
|||
class Trainer:
|
||||
"""模型训练器，负责整个训练流程的管理"""
|
||||
|
||||
def __init__(self, model, loss, optimizer,
|
||||
train_loader, val_loader, test_loader, scaler,
|
||||
args, lr_scheduler=None,):
|
||||
def __init__(self, model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args, lr_scheduler=None):
|
||||
# 设备和基本参数
|
||||
self.config = args
|
||||
self.device = args["basic"]["device"]
|
||||
train_args = args["train"]
|
||||
self.args = args["train"]
|
||||
|
||||
# 模型和训练相关组件
|
||||
self.model = model
|
||||
self.loss = loss
|
||||
self.optimizer = optimizer
|
||||
self.lr_scheduler = lr_scheduler
|
||||
self.model, self.loss, self.optimizer, self.lr_scheduler = model, loss, optimizer, lr_scheduler
|
||||
|
||||
# 数据加载器
|
||||
self.train_loader = train_loader
|
||||
self.val_loader = val_loader
|
||||
self.test_loader = test_loader
|
||||
self.train_loader, self.val_loader, self.test_loader = train_loader, val_loader, test_loader
|
||||
|
||||
# 数据处理工具
|
||||
self.scaler = scaler
|
||||
self.args = train_args
|
||||
|
||||
# 初始化路径、日志和统计
|
||||
self._initialize_paths(train_args)
|
||||
self._initialize_logger(train_args)
|
||||
self._initialize_paths(self.args)
|
||||
self._initialize_logger(self.args)
|
||||
|
||||
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")
|
||||
log_dir = args["log_dir"]
|
||||
self.best_path = os.path.join(log_dir, "best_model.pth")
|
||||
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):
|
||||
"""初始化日志记录器"""
|
||||
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']}")
|
||||
log_dir = args["log_dir"]
|
||||
if not args["debug"]:
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
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):
|
||||
"""运行一个训练/验证/测试epoch"""
|
||||
# 设置模型模式和是否进行优化
|
||||
if mode == "train": self.model.train(); optimizer_step = True
|
||||
else: self.model.eval(); optimizer_step = False
|
||||
self.model.train() if mode == "train" else self.model.eval()
|
||||
optimizer_step = mode == "train"
|
||||
|
||||
# 初始化变量
|
||||
total_loss = 0
|
||||
|
|
@ -60,105 +58,111 @@ class Trainer:
|
|||
# 训练/验证循环
|
||||
with torch.set_grad_enabled(optimizer_step):
|
||||
progress_bar = tqdm(
|
||||
enumerate(dataloader),
|
||||
dataloader,
|
||||
total=len(dataloader),
|
||||
desc=f"{mode.capitalize()} Epoch {epoch}"
|
||||
)
|
||||
for _, (data, target) in progress_bar:
|
||||
# 转移数据
|
||||
data = data.to(self.device)
|
||||
target = target.to(self.device)
|
||||
for data, target in progress_bar:
|
||||
# 转移数据并提取标签
|
||||
data, target = data.to(self.device), target.to(self.device)
|
||||
label = target[..., : self.args["output_dim"]]
|
||||
# 计算loss和反归一化loss
|
||||
|
||||
# 计算输出
|
||||
output = self.model(data)
|
||||
|
||||
# 我的调试开关
|
||||
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
|
||||
|
||||
# 计算损失
|
||||
loss = self.loss(output, label)
|
||||
d_output = self.scaler.inverse_transform(output)
|
||||
d_label = self.scaler.inverse_transform(label)
|
||||
d_loss = self.loss(d_output, d_label)
|
||||
|
||||
# 累积损失和预测结果
|
||||
total_loss += d_loss.item()
|
||||
y_pred.append(d_output.detach().cpu())
|
||||
y_true.append(d_label.detach().cpu())
|
||||
|
||||
# 反向传播和优化（仅在训练模式）
|
||||
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()
|
||||
|
||||
# 更新进度条
|
||||
progress_bar.set_postfix(loss=d_loss.item())
|
||||
|
||||
# 合并所有批次的预测结果
|
||||
y_pred = torch.cat(y_pred, dim=0)
|
||||
y_true = torch.cat(y_true, dim=0)
|
||||
y_pred, y_true = torch.cat(y_pred, dim=0), 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"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"
|
||||
)
|
||||
|
||||
return avg_loss
|
||||
|
||||
def train_epoch(self, epoch):
|
||||
return self._run_epoch(epoch, self.train_loader, "train")
|
||||
|
||||
def val_epoch(self, epoch):
|
||||
return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
|
||||
|
||||
def test_epoch(self, epoch):
|
||||
return self._run_epoch(epoch, self.test_loader, "test")
|
||||
|
||||
def train(self):
|
||||
# 初始化记录
|
||||
best_model, best_test_model = None, None
|
||||
best_loss, best_test_loss = float("inf"), float("inf")
|
||||
best_model = best_test_model = None
|
||||
best_loss = best_test_loss = float("inf")
|
||||
not_improved_count = 0
|
||||
|
||||
# 开始训练
|
||||
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)
|
||||
train_epoch_loss = self._run_epoch(epoch, self.train_loader, "train")
|
||||
val_epoch_loss = self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
|
||||
test_epoch_loss = self._run_epoch(epoch, self.test_loader, "test")
|
||||
|
||||
# 检查梯度爆炸
|
||||
if train_epoch_loss > 1e6:
|
||||
self.logger.warning("Gradient explosion detected. Ending...")
|
||||
break
|
||||
|
||||
# 更新最佳验证模型
|
||||
if val_epoch_loss < best_loss:
|
||||
best_loss = val_epoch_loss
|
||||
not_improved_count = 0
|
||||
best_loss, not_improved_count = val_epoch_loss, 0
|
||||
best_model = copy.deepcopy(self.model.state_dict())
|
||||
self.logger.info("Best validation model saved!")
|
||||
else:
|
||||
not_improved_count += 1
|
||||
# 早停
|
||||
|
||||
# 早停检查
|
||||
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"]:
|
||||
self._save_best_models(best_model, best_test_model)
|
||||
|
||||
# 最终评估
|
||||
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"]
|
||||
):
|
||||
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."
|
||||
)
|
||||
|
|
@ -190,58 +194,43 @@ class Trainer:
|
|||
@staticmethod
|
||||
def test(model, args, data_loader, scaler, logger, path=None):
|
||||
"""对模型进行评估并输出性能指标"""
|
||||
# 确定设备信息
|
||||
device = None
|
||||
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:
|
||||
# 验证参数类型
|
||||
if not isinstance(args, dict):
|
||||
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:
|
||||
checkpoint = torch.load(path)
|
||||
model.load_state_dict(checkpoint["state_dict"])
|
||||
model.to(device)
|
||||
|
||||
# 设置为评估模式
|
||||
# 设置为评估模式并收集预测结果
|
||||
model.eval()
|
||||
|
||||
# 收集预测和真实标签
|
||||
y_pred, y_true = [], []
|
||||
|
||||
# 不计算梯度的情况下进行预测
|
||||
with torch.no_grad():
|
||||
for data, target in data_loader:
|
||||
# 将数据和标签移动到指定设备
|
||||
data = data.to(device)
|
||||
target = target.to(device)
|
||||
|
||||
data, target = data.to(device), target.to(device)
|
||||
label = target[..., : output_dim]
|
||||
|
||||
output = model(data)
|
||||
y_pred.append(output.detach().cpu())
|
||||
y_true.append(label.detach().cpu())
|
||||
|
||||
# 反归一化并计算指标
|
||||
d_y_pred = scaler.inverse_transform(torch.cat(y_pred, 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]):
|
||||
|
|
@ -254,9 +243,5 @@ class Trainer:
|
|||
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)
|
||||
logger.info( f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}")
|
||||
|
||||
@staticmethod
|
||||
def _compute_sampling_threshold(global_step, k):
|
||||
return k / (k + math.exp(global_step / k))
|
||||
avg_mae, avg_rmse, avg_mape = all_metrics(d_y_pred, d_y_true, mae_thresh, mape_thresh)
|
||||
logger.info(f"Average Horizon, MAE: {avg_mae:.4f}, RMSE: {avg_rmse:.4f}, MAPE: {avg_mape:.4f}")
|
||||
|
|
|
|||
Loading…
Reference in New Issue