mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-21 21:22:05 +08:00
560 lines
19 KiB
Python
560 lines
19 KiB
Python
"""带类权重的训练器：处理分类任务的数据不均衡问题"""
|
||
|
||
from typing import Dict, List, Optional, Tuple
|
||
from dataclasses import dataclass, field
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import DataLoader
|
||
from tqdm import tqdm
|
||
|
||
|
||
@dataclass
|
||
class ClassWeights:
|
||
"""分类任务的类权重"""
|
||
pdi: Optional[torch.Tensor] = None # [4] for 4 PDI classes
|
||
ee: Optional[torch.Tensor] = None # [3] for 3 EE classes
|
||
toxic: Optional[torch.Tensor] = None # [2] for binary toxic
|
||
|
||
|
||
@dataclass
|
||
class LossWeightsBalanced:
|
||
"""各任务的损失权重（与 trainer.py 的 LossWeights 兼容）"""
|
||
size: float = 1.0
|
||
pdi: float = 1.0
|
||
ee: float = 1.0
|
||
delivery: float = 1.0
|
||
biodist: float = 1.0
|
||
toxic: float = 1.0
|
||
moe_lb: float = 0.01 # MoE load-balancing 系数（仅在 use_moe=True 时生效）
|
||
|
||
|
||
def compute_class_weights_from_loader(
|
||
loader: DataLoader,
|
||
n_pdi_classes: int = 2,
|
||
n_ee_classes: int = 3,
|
||
n_toxic_classes: int = 2,
|
||
smoothing: float = 0.1,
|
||
) -> ClassWeights:
|
||
"""
|
||
从 DataLoader 统计类别频次并计算类权重。
|
||
|
||
使用 inverse frequency 方式：weight_c = N / (n_classes * count_c)
|
||
加 smoothing 避免极端权重。
|
||
|
||
Args:
|
||
loader: DataLoader（需要遍历一次）
|
||
n_pdi_classes: PDI 类别数
|
||
n_ee_classes: EE 类别数
|
||
n_toxic_classes: toxic 类别数
|
||
smoothing: 平滑系数（防止除零和极端权重）
|
||
|
||
Returns:
|
||
ClassWeights 对象，包含各分类任务的类权重张量
|
||
"""
|
||
pdi_counts = torch.zeros(n_pdi_classes)
|
||
ee_counts = torch.zeros(n_ee_classes)
|
||
toxic_counts = torch.zeros(n_toxic_classes)
|
||
|
||
for batch in loader:
|
||
targets = batch["targets"]
|
||
mask = batch["mask"]
|
||
|
||
# PDI
|
||
if "pdi" in targets and "pdi" in mask:
|
||
m = mask["pdi"]
|
||
if m.any():
|
||
labels = targets["pdi"][m]
|
||
for c in range(n_pdi_classes):
|
||
pdi_counts[c] += (labels == c).sum().item()
|
||
|
||
# EE
|
||
if "ee" in targets and "ee" in mask:
|
||
m = mask["ee"]
|
||
if m.any():
|
||
labels = targets["ee"][m]
|
||
for c in range(n_ee_classes):
|
||
ee_counts[c] += (labels == c).sum().item()
|
||
|
||
# Toxic
|
||
if "toxic" in targets and "toxic" in mask:
|
||
m = mask["toxic"]
|
||
if m.any():
|
||
labels = targets["toxic"][m]
|
||
for c in range(n_toxic_classes):
|
||
toxic_counts[c] += (labels == c).sum().item()
|
||
|
||
def counts_to_weights(counts: torch.Tensor, n_classes: int) -> Optional[torch.Tensor]:
|
||
"""将计数转换为类权重"""
|
||
total = counts.sum().item()
|
||
if total == 0:
|
||
return None
|
||
# Inverse frequency with smoothing
|
||
counts = counts + smoothing
|
||
weights = total / (n_classes * counts)
|
||
# Normalize to mean=1
|
||
weights = weights / weights.mean()
|
||
return weights
|
||
|
||
return ClassWeights(
|
||
pdi=counts_to_weights(pdi_counts, n_pdi_classes),
|
||
ee=counts_to_weights(ee_counts, n_ee_classes),
|
||
toxic=counts_to_weights(toxic_counts, n_toxic_classes),
|
||
)
|
||
|
||
|
||
def compute_multitask_loss_balanced(
|
||
outputs: Dict[str, torch.Tensor],
|
||
targets: Dict[str, torch.Tensor],
|
||
mask: Dict[str, torch.Tensor],
|
||
task_weights: Optional[LossWeightsBalanced] = None,
|
||
class_weights: Optional[ClassWeights] = None,
|
||
model: Optional[nn.Module] = None,
|
||
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||
"""
|
||
计算带类权重的多任务损失。
|
||
|
||
Args:
|
||
outputs: 模型输出
|
||
targets: 真实标签
|
||
mask: 有效样本掩码
|
||
task_weights: 各任务权重
|
||
class_weights: 分类任务的类权重
|
||
|
||
Returns:
|
||
(total_loss, loss_dict) 总损失和各任务损失
|
||
"""
|
||
task_weights = task_weights or LossWeightsBalanced()
|
||
class_weights = class_weights or ClassWeights()
|
||
|
||
# 若模型 head 暴露 log_vars，则启用不确定性加权（自动平衡），否则回退到标量权重
|
||
log_vars = None
|
||
if model is not None and hasattr(model, "head") and hasattr(model.head, "log_vars"):
|
||
log_vars = model.head.log_vars
|
||
|
||
losses = {}
|
||
device = next(iter(outputs.values())).device
|
||
total_loss = torch.tensor(0.0, device=device)
|
||
|
||
def _add(name: str, raw_loss: torch.Tensor, scalar_w: float) -> None:
|
||
nonlocal total_loss
|
||
losses[name] = raw_loss
|
||
if log_vars is not None and name in log_vars:
|
||
s = log_vars[name]
|
||
total_loss = total_loss + torch.exp(-s) * raw_loss + 0.5 * s
|
||
else:
|
||
total_loss = total_loss + scalar_w * raw_loss
|
||
|
||
# size: MSE
|
||
if "size" in targets and mask["size"].any():
|
||
m = mask["size"]
|
||
_add("size", F.mse_loss(outputs["size"][m].squeeze(-1), targets["size"][m]), task_weights.size)
|
||
|
||
# delivery: MSE
|
||
if "delivery" in targets and mask["delivery"].any():
|
||
m = mask["delivery"]
|
||
_add("delivery", F.mse_loss(outputs["delivery"][m].squeeze(-1), targets["delivery"][m]), task_weights.delivery)
|
||
|
||
# pdi: 类加权 CE
|
||
if "pdi" in targets and mask["pdi"].any():
|
||
m = mask["pdi"]
|
||
w = class_weights.pdi.to(device) if class_weights.pdi is not None else None
|
||
_add("pdi", F.cross_entropy(outputs["pdi"][m], targets["pdi"][m], weight=w), task_weights.pdi)
|
||
|
||
# ee: 类加权 CE
|
||
if "ee" in targets and mask["ee"].any():
|
||
m = mask["ee"]
|
||
w = class_weights.ee.to(device) if class_weights.ee is not None else None
|
||
_add("ee", F.cross_entropy(outputs["ee"][m], targets["ee"][m], weight=w), task_weights.ee)
|
||
|
||
# toxic: 类加权 CE
|
||
if "toxic" in targets and mask["toxic"].any():
|
||
m = mask["toxic"]
|
||
w = class_weights.toxic.to(device) if class_weights.toxic is not None else None
|
||
_add("toxic", F.cross_entropy(outputs["toxic"][m], targets["toxic"][m], weight=w), task_weights.toxic)
|
||
|
||
# biodist: KL
|
||
if "biodist" in targets and mask["biodist"].any():
|
||
m = mask["biodist"]
|
||
kl = F.kl_div(outputs["biodist"][m].log().clamp(min=-100), targets["biodist"][m], reduction="batchmean")
|
||
_add("biodist", kl, task_weights.biodist)
|
||
|
||
# MoE load-balancing aux loss（固定小权重，不纳入不确定性加权）
|
||
if model is not None and hasattr(model, "get_last_moe_extras"):
|
||
extras = model.get_last_moe_extras()
|
||
if extras is not None and "lb_loss" in extras:
|
||
losses["moe_lb"] = extras["lb_loss"]
|
||
total_loss = total_loss + task_weights.moe_lb * losses["moe_lb"]
|
||
|
||
return total_loss, losses
|
||
|
||
|
||
def train_epoch_balanced(
|
||
model: nn.Module,
|
||
loader: DataLoader,
|
||
optimizer: torch.optim.Optimizer,
|
||
device: torch.device,
|
||
task_weights: Optional[LossWeightsBalanced] = None,
|
||
class_weights: Optional[ClassWeights] = None,
|
||
) -> Dict[str, float]:
|
||
"""带类权重的训练一个 epoch"""
|
||
model.train()
|
||
total_loss = 0.0
|
||
task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
|
||
n_batches = 0
|
||
|
||
for batch in tqdm(loader, desc="Training", leave=False):
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
|
||
targets = {k: v.to(device) for k, v in batch["targets"].items()}
|
||
mask = {k: v.to(device) for k, v in batch["mask"].items()}
|
||
|
||
optimizer.zero_grad()
|
||
outputs = model(smiles, tabular)
|
||
loss, losses = compute_multitask_loss_balanced(
|
||
outputs, targets, mask, task_weights, class_weights, model=model,
|
||
)
|
||
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||
optimizer.step()
|
||
|
||
total_loss += loss.item()
|
||
for k, v in losses.items():
|
||
task_losses[k] += v.item()
|
||
n_batches += 1
|
||
|
||
return {
|
||
"loss": total_loss / n_batches,
|
||
**{f"loss_{k}": v / n_batches for k, v in task_losses.items() if v > 0},
|
||
}
|
||
|
||
|
||
@torch.no_grad()
|
||
def validate_balanced(
|
||
model: nn.Module,
|
||
loader: DataLoader,
|
||
device: torch.device,
|
||
task_weights: Optional[LossWeightsBalanced] = None,
|
||
class_weights: Optional[ClassWeights] = None,
|
||
) -> Dict[str, float]:
|
||
"""带类权重的验证"""
|
||
model.eval()
|
||
total_loss = 0.0
|
||
task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
|
||
n_batches = 0
|
||
|
||
# 用于计算准确率
|
||
correct = {k: 0 for k in ["pdi", "ee", "toxic"]}
|
||
total = {k: 0 for k in ["pdi", "ee", "toxic"]}
|
||
|
||
for batch in tqdm(loader, desc="Validating", leave=False):
|
||
smiles = batch["smiles"]
|
||
tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
|
||
targets = {k: v.to(device) for k, v in batch["targets"].items()}
|
||
mask = {k: v.to(device) for k, v in batch["mask"].items()}
|
||
|
||
outputs = model(smiles, tabular)
|
||
loss, losses = compute_multitask_loss_balanced(
|
||
outputs, targets, mask, task_weights, class_weights, model=model,
|
||
)
|
||
|
||
total_loss += loss.item()
|
||
for k, v in losses.items():
|
||
task_losses[k] += v.item()
|
||
n_batches += 1
|
||
|
||
# 计算分类准确率
|
||
for k in ["pdi", "ee", "toxic"]:
|
||
if k in targets and mask[k].any():
|
||
m = mask[k]
|
||
pred = outputs[k][m].argmax(dim=-1)
|
||
tgt = targets[k][m]
|
||
correct[k] += (pred == tgt).sum().item()
|
||
total[k] += m.sum().item()
|
||
|
||
metrics = {
|
||
"loss": total_loss / n_batches,
|
||
**{f"loss_{k}": v / n_batches for k, v in task_losses.items() if v > 0},
|
||
}
|
||
|
||
# 添加准确率
|
||
for k in ["pdi", "ee", "toxic"]:
|
||
if total[k] > 0:
|
||
metrics[f"acc_{k}"] = correct[k] / total[k]
|
||
|
||
return metrics
|
||
|
||
|
||
BACKBONE_PREFIXES = (
|
||
"token_projector.",
|
||
"set_transformer.",
|
||
"fusion.",
|
||
"moe.",
|
||
"llm_prompt.",
|
||
)
|
||
|
||
FROM_SCRATCH_PREFIXES = (
|
||
"moe.",
|
||
"llm_prompt.",
|
||
"fusion.g_moe",
|
||
"fusion.g_llm",
|
||
)
|
||
|
||
|
||
def build_optimizer(
|
||
model: nn.Module,
|
||
lr: float,
|
||
weight_decay: float,
|
||
backbone_lr_ratio: float = 1.0,
|
||
size_wd_mult: float = 5.0,
|
||
) -> torch.optim.AdamW:
|
||
"""
|
||
构建 AdamW 优化器，支持分层学习率。
|
||
|
||
仅收集 requires_grad=True 的参数（冻结的 MolT5 encoder 被排除）。
|
||
当 backbone_lr_ratio < 1.0 时，backbone 参数使用 lr * backbone_lr_ratio，
|
||
其余参数（task heads 等）使用 lr。backbone_lr_ratio = 1.0 等价于统一学习率。
|
||
size_head 信号弱、易过拟合 -> 单独施加 size_wd_mult 倍 weight_decay。
|
||
"""
|
||
trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
|
||
size_wd = weight_decay * size_wd_mult
|
||
|
||
def is_size(name: str) -> bool:
|
||
return "size_head" in name
|
||
|
||
size_params = [p for n, p in trainable if is_size(n)]
|
||
|
||
if backbone_lr_ratio >= 1.0:
|
||
other = [p for n, p in trainable if not is_size(n)]
|
||
return torch.optim.AdamW(
|
||
[
|
||
{"params": other, "weight_decay": weight_decay},
|
||
{"params": size_params, "weight_decay": size_wd},
|
||
],
|
||
lr=lr,
|
||
)
|
||
|
||
backbone_params, head_params = [], []
|
||
for name, param in trainable:
|
||
if is_size(name):
|
||
continue
|
||
if name.startswith(BACKBONE_PREFIXES) and not name.startswith(FROM_SCRATCH_PREFIXES):
|
||
backbone_params.append(param)
|
||
else:
|
||
head_params.append(param)
|
||
|
||
return torch.optim.AdamW(
|
||
[
|
||
{"params": backbone_params, "lr": lr * backbone_lr_ratio, "weight_decay": weight_decay},
|
||
{"params": head_params, "lr": lr, "weight_decay": weight_decay},
|
||
{"params": size_params, "lr": lr, "weight_decay": size_wd},
|
||
],
|
||
)
|
||
|
||
|
||
class EarlyStoppingBalanced:
|
||
"""早停机制（与 trainer.py 的 EarlyStopping 兼容）"""
|
||
|
||
def __init__(self, patience: int = 10, min_delta: float = 0.0):
|
||
self.patience = patience
|
||
self.min_delta = min_delta
|
||
self.counter = 0
|
||
self.best_loss = float("inf")
|
||
self.best_epoch = 0
|
||
self.should_stop = False
|
||
|
||
def __call__(self, val_loss: float, epoch: int = 0) -> bool:
|
||
if val_loss < self.best_loss - self.min_delta:
|
||
self.best_loss = val_loss
|
||
self.best_epoch = epoch
|
||
self.counter = 0
|
||
else:
|
||
self.counter += 1
|
||
if self.counter >= self.patience:
|
||
self.should_stop = True
|
||
return self.should_stop
|
||
|
||
def get_best_epoch(self) -> int:
|
||
"""获取最佳 epoch（1-indexed）"""
|
||
return self.best_epoch + 1
|
||
|
||
|
||
def train_with_early_stopping(
|
||
model: nn.Module,
|
||
train_loader: DataLoader,
|
||
val_loader: DataLoader,
|
||
device: torch.device,
|
||
lr: float = 1e-4,
|
||
weight_decay: float = 1e-5,
|
||
epochs: int = 100,
|
||
patience: int = 15,
|
||
task_weights: Optional[LossWeightsBalanced] = None,
|
||
class_weights: Optional[ClassWeights] = None,
|
||
backbone_lr_ratio: float = 1.0,
|
||
) -> Dict:
|
||
"""
|
||
带早停的完整训练流程。
|
||
|
||
Returns:
|
||
Dict with keys: history, best_val_loss, best_epoch, best_state
|
||
"""
|
||
model = model.to(device)
|
||
optimizer = build_optimizer(model, lr, weight_decay, backbone_lr_ratio)
|
||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||
optimizer, mode="min", factor=0.5, patience=5
|
||
)
|
||
early_stopping = EarlyStoppingBalanced(patience=patience, min_delta=1e-3)
|
||
|
||
history = {"train": [], "val": []}
|
||
best_val_loss = float("inf")
|
||
best_state = None
|
||
|
||
for epoch in range(epochs):
|
||
# Train
|
||
train_metrics = train_epoch_balanced(
|
||
model, train_loader, optimizer, device, task_weights, class_weights
|
||
)
|
||
|
||
# Validate
|
||
val_metrics = validate_balanced(
|
||
model, val_loader, device, task_weights, class_weights
|
||
)
|
||
|
||
history["train"].append(train_metrics)
|
||
history["val"].append(val_metrics)
|
||
|
||
# Learning rate scheduling
|
||
scheduler.step(val_metrics["loss"])
|
||
|
||
# Save best model
|
||
if val_metrics["loss"] < best_val_loss:
|
||
best_val_loss = val_metrics["loss"]
|
||
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
|
||
|
||
# Early stopping
|
||
if early_stopping(val_metrics["loss"], epoch):
|
||
break
|
||
|
||
# Restore best model（QLoRA 4-bit 基座的量化元数据键用 strict=False 忽略）
|
||
if best_state is not None:
|
||
model.load_state_dict(best_state, strict=False)
|
||
|
||
return {
|
||
"history": history,
|
||
"best_val_loss": best_val_loss,
|
||
"best_epoch": early_stopping.get_best_epoch(),
|
||
"best_state": best_state,
|
||
"epochs_trained": len(history["train"]),
|
||
}
|
||
|
||
|
||
def train_fixed_epochs(
|
||
model: nn.Module,
|
||
train_loader: DataLoader,
|
||
val_loader: Optional[DataLoader],
|
||
device: torch.device,
|
||
lr: float = 1e-4,
|
||
weight_decay: float = 1e-5,
|
||
epochs: int = 50,
|
||
task_weights: Optional[LossWeightsBalanced] = None,
|
||
class_weights: Optional[ClassWeights] = None,
|
||
use_cosine_annealing: bool = True,
|
||
use_swa: bool = False,
|
||
swa_start_epoch: Optional[int] = None,
|
||
backbone_lr_ratio: float = 1.0,
|
||
freeze_backbone_epochs: int = 0,
|
||
) -> Dict:
|
||
"""
|
||
固定 epoch 数的训练（不使用 early stopping）。
|
||
|
||
用于外层 CV 训练和最终训练。
|
||
|
||
Args:
|
||
model: 模型
|
||
train_loader: 训练数据
|
||
val_loader: 验证数据（可选，仅用于监控）
|
||
device: 设备
|
||
lr: 学习率
|
||
weight_decay: 权重衰减
|
||
epochs: 训练轮数
|
||
task_weights: 任务权重
|
||
class_weights: 类权重
|
||
use_cosine_annealing: 是否使用 CosineAnnealingLR
|
||
use_swa: 是否使用 SWA
|
||
swa_start_epoch: SWA 开始的 epoch（默认为 epochs * 0.75）
|
||
backbone_lr_ratio: backbone 学习率相对于 head 的比例（1.0 = 统一学习率）
|
||
|
||
Returns:
|
||
Dict with keys: history, final_state
|
||
"""
|
||
model = model.to(device)
|
||
optimizer = build_optimizer(model, lr, weight_decay, backbone_lr_ratio)
|
||
|
||
backbone_named = [
|
||
(n, p) for n, p in model.named_parameters()
|
||
if n.startswith(BACKBONE_PREFIXES) and not n.startswith(FROM_SCRATCH_PREFIXES)
|
||
and p.requires_grad
|
||
]
|
||
if freeze_backbone_epochs > 0:
|
||
for _, p in backbone_named:
|
||
p.requires_grad_(False)
|
||
|
||
if use_cosine_annealing:
|
||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
|
||
else:
|
||
scheduler = None
|
||
|
||
# SWA setup
|
||
swa_model = None
|
||
swa_scheduler = None
|
||
if use_swa:
|
||
from torch.optim.swa_utils import AveragedModel, SWALR
|
||
swa_model = AveragedModel(model)
|
||
swa_start = swa_start_epoch or int(epochs * 0.75)
|
||
swa_scheduler = SWALR(optimizer, swa_lr=lr * 0.1)
|
||
|
||
history = {"train": [], "val": []}
|
||
|
||
for epoch in range(epochs):
|
||
if freeze_backbone_epochs > 0 and epoch == freeze_backbone_epochs:
|
||
for _, p in backbone_named:
|
||
p.requires_grad_(True)
|
||
# Train
|
||
train_metrics = train_epoch_balanced(
|
||
model, train_loader, optimizer, device, task_weights, class_weights
|
||
)
|
||
history["train"].append(train_metrics)
|
||
|
||
# Validate (optional)
|
||
if val_loader is not None:
|
||
val_metrics = validate_balanced(
|
||
model, val_loader, device, task_weights, class_weights
|
||
)
|
||
history["val"].append(val_metrics)
|
||
|
||
# Scheduler step
|
||
if use_swa and epoch >= swa_start:
|
||
swa_model.update_parameters(model)
|
||
swa_scheduler.step()
|
||
elif scheduler is not None:
|
||
scheduler.step()
|
||
|
||
# Finalize SWA
|
||
final_state = None
|
||
if use_swa and swa_model is not None:
|
||
# Update batch normalization statistics
|
||
torch.optim.swa_utils.update_bn(train_loader, swa_model, device=device)
|
||
final_state = {k: v.cpu().clone() for k, v in swa_model.module.state_dict().items()}
|
||
else:
|
||
final_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
|
||
|
||
return {
|
||
"history": history,
|
||
"final_state": final_state,
|
||
"epochs_trained": epochs,
|
||
}
|
||
|