393 lines
15 KiB
Python
393 lines
15 KiB
Python
import math
|
|
import os
|
|
import time
|
|
import copy
|
|
import torch.nn.functional as F
|
|
import torch
|
|
from torch import nn
|
|
|
|
from tqdm import tqdm
|
|
from utils.logger import get_logger
|
|
from utils.loss_function import all_metrics
|
|
from model.STMLP.STMLP import STMLP
|
|
from utils.training_stats import TrainingStats
|
|
|
|
|
|
class Trainer:
|
|
def __init__(
|
|
self,
|
|
model,
|
|
loss,
|
|
optimizer,
|
|
train_loader,
|
|
val_loader,
|
|
test_loader,
|
|
scaler,
|
|
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 = train_args
|
|
|
|
# 统计信息
|
|
self.train_per_epoch = len(train_loader)
|
|
self.val_per_epoch = len(val_loader) if val_loader else 0
|
|
|
|
# 初始化路径、日志和统计
|
|
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:
|
|
self.logger.info(
|
|
f"当前使用预训练模式，预训练后请移动教师模型到"
|
|
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
|
|
else:
|
|
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"]]
|
|
|
|
if self.args["teacher_stu"]:
|
|
# 教师-学生蒸馏模式
|
|
output, out_, _ = self.model(data)
|
|
gout, tout, sout = self.tmodel(data)
|
|
|
|
# 计算原始loss
|
|
loss1 = self.loss(output, label)
|
|
|
|
# 计算蒸馏相关loss
|
|
scl = self.loss_cls(out_, sout)
|
|
kl_loss = nn.KLDivLoss(
|
|
reduction="batchmean", log_target=True
|
|
).cuda()
|
|
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
|
|
|
|
# 检查output和label的shape是否一致
|
|
if output.shape == label.shape:
|
|
print(f"✓ Test passed: output shape {output.shape} matches label shape {label.shape}")
|
|
import sys
|
|
sys.exit(0)
|
|
else:
|
|
print(f"✗ Test failed: output shape {output.shape} does not match label shape {label.shape}")
|
|
import sys
|
|
sys.exit(1)
|
|
else:
|
|
# 普通训练模式
|
|
output, out_, _ = self.model(data)
|
|
loss = self.loss(output, label)
|
|
|
|
# 检查output和label的shape是否一致
|
|
if output.shape == label.shape:
|
|
print(f"✓ Test passed: output shape {output.shape} matches label shape {label.shape}")
|
|
import sys
|
|
sys.exit(0)
|
|
else:
|
|
print(f"✗ Test failed: output shape {output.shape} does not match label shape {label.shape}")
|
|
import sys
|
|
sys.exit(1)
|
|
|
|
# 反归一化
|
|
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}"
|
|
)
|
|
|
|
# 更新 tqdm 的进度
|
|
pbar.update(1)
|
|
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"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):
|
|
return self._run_epoch(epoch, self.train_loader, "train")
|
|
|
|
def val_epoch(self, epoch):
|
|
return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")
|
|
|
|
def test_epoch(self, epoch):
|
|
return self._run_epoch(epoch, self.test_loader, "test")
|
|
|
|
def train(self):
|
|
"""执行完整的训练流程"""
|
|
best_model, best_test_model = None, None
|
|
best_loss, best_test_loss = float("inf"), float("inf")
|
|
not_improved_count = 0
|
|
|
|
self.stats.start_training()
|
|
self.logger.info("Training process started")
|
|
for epoch in range(1, self.args["epochs"] + 1):
|
|
train_epoch_loss = self.train_epoch(epoch)
|
|
val_epoch_loss = self.val_epoch(epoch)
|
|
test_epoch_loss = self.test_epoch(epoch)
|
|
|
|
if train_epoch_loss > 1e6:
|
|
self.logger.warning("Gradient explosion detected. Ending...")
|
|
break
|
|
|
|
if val_epoch_loss < best_loss:
|
|
best_loss = val_epoch_loss
|
|
not_improved_count = 0
|
|
best_model = copy.deepcopy(self.model.state_dict())
|
|
torch.save(best_model, self.best_path)
|
|
torch.save(best_model, self.pretrain_path)
|
|
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())
|
|
torch.save(best_test_model, self.best_test_path)
|
|
torch.save(best_model, self.pretrain_best_path)
|
|
|
|
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.stats.end_training()
|
|
self.stats.report(self.logger)
|
|
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)
|
|
self.logger.info("Testing on best validation model")
|
|
self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
|
|
|
|
self.model.load_state_dict(best_test_model)
|
|
self.logger.info("Testing on best test model")
|
|
self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)
|
|
|
|
def loadTeacher(self, args):
|
|
model_path = (
|
|
f"./pre-train/{args['model']['type']}/{args['data']['type']}/best_model.pth"
|
|
)
|
|
try:
|
|
# 尝试加载教师模型权重
|
|
state_dict = torch.load(model_path)
|
|
self.logger.info(f"成功加载教师模型权重: {model_path}")
|
|
|
|
# 初始化并返回教师模型
|
|
args["model"]["model_type"] = "teacher"
|
|
tmodel = STMLP(args["model"])
|
|
tmodel = tmodel.to(args["device"])
|
|
tmodel.load_state_dict(state_dict, strict=False)
|
|
return tmodel
|
|
|
|
except FileNotFoundError:
|
|
# 如果找不到权重文件，记录日志并修改 args
|
|
self.logger.error(
|
|
f"未找到教师模型权重文件: {model_path}。切换到预训练模式训练老师权重。\n"
|
|
f"在预训练完成后，再次启动模型则为蒸馏模式"
|
|
)
|
|
self.args["teacher_stu"] = False
|
|
return None
|
|
|
|
def loss_cls(self, x1, x2):
|
|
temperature = 0.05
|
|
x1 = F.normalize(x1, p=2, dim=-1)
|
|
x2 = F.normalize(x2, p=2, dim=-1)
|
|
weight = F.cosine_similarity(x1, x2, dim=-1)
|
|
batch_size = x1.size()[0]
|
|
# neg score
|
|
out = torch.cat([x1, x2], dim=0)
|
|
neg = torch.exp(
|
|
torch.matmul(out, out.transpose(2, 3).contiguous()) / temperature
|
|
)
|
|
|
|
pos = torch.exp(torch.sum(x1 * x2, dim=-1) * weight / temperature)
|
|
# pos = torch.exp(torch.sum(x1 * x2, dim=-1) / temperature)
|
|
pos = torch.cat([pos, pos], dim=0).sum(dim=1)
|
|
|
|
Ng = neg.sum(dim=-1).sum(dim=1)
|
|
|
|
loss = (-torch.log(pos / (pos + Ng))).mean()
|
|
|
|
return loss
|
|
|
|
@staticmethod
|
|
def test(model, args, data_loader, scaler, logger, path=None):
|
|
"""对模型进行评估并输出性能指标"""
|
|
# 加载模型检查点（如果提供了路径）
|
|
if path:
|
|
checkpoint = torch.load(path)
|
|
model.load_state_dict(checkpoint["state_dict"])
|
|
model.to(args["basic"]["device"])
|
|
|
|
# 设置为评估模式
|
|
model.eval()
|
|
|
|
# 收集预测和真实标签
|
|
y_pred, y_true = [], []
|
|
|
|
# 不计算梯度的情况下进行预测
|
|
with torch.no_grad():
|
|
for data, target in data_loader:
|
|
label = target[..., : args["output_dim"]]
|
|
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))
|
|
|
|
# 计算并记录每个时间步的指标
|
|
for t in range(d_y_true.shape[1]):
|
|
mae, rmse, mape = all_metrics(
|
|
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}")
|
|
|
|
# 计算并记录平均指标
|
|
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):
|
|
return k / (k + math.exp(global_step / k))
|