mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-21 21:22:05 +08:00
feat: 回归旁路（regression bypass）+ --reg-bypass 开关
- fusion.py: ResidualConcatFusion 额外返回 pooled（纯 chem+tab 注意力池化） - models.py: _backbone_from_stacked 存 _last_pooled；forward 按 reg_bypass 决定是否传入 - heads.py: MultiTaskHead 接收 x_numeric，回归头(size/delivery)走纯数值，分类/biodist 走 fused - nested_cv_optuna.py: 新增 --reg-bypass 开关（typer），经 llm_kwargs 传入 fix: LNPModelWithoutMPNN 未将 reg_bypass 转发给 super().__init__()， 导致开关静默失效、on/off 产生完全相同结果（已验证修复：off→False, on→True）
This commit is contained in:
DicongLi
parent
0c6828076b
commit
104dfef94c
@ -1,125 +1,126 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class RegressionHead(nn.Module):
|
||||
"""回归任务头：输出单个 float 值"""
|
||||
|
||||
def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, 1]"""
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class ClassificationHead(nn.Module):
|
||||
"""分类任务头：输出 logits"""
|
||||
|
||||
def __init__(
|
||||
self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, num_classes),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, num_classes] (logits)"""
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class DistributionHead(nn.Module):
|
||||
"""分布任务头：输出和为 1 的概率分布（用于 Biodistribution）"""
|
||||
|
||||
def __init__(
|
||||
self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, num_outputs),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, num_outputs] (softmax, sum=1)"""
|
||||
logits = self.net(x)
|
||||
return F.softmax(logits, dim=-1)
|
||||
|
||||
|
||||
class MultiTaskHead(nn.Module):
|
||||
"""
|
||||
多任务预测头，根据任务配置自动创建对应的子头。
|
||||
|
||||
输出:
|
||||
- size: [B, 1] 回归
|
||||
- pdi: [B, 2] 分类 logits
|
||||
- ee: [B, 3] 分类 logits
|
||||
- delivery: [B, 1] 回归
|
||||
- biodist: [B, 7] softmax 分布
|
||||
- toxic: [B, 2] 二分类 logits
|
||||
"""
|
||||
|
||||
def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
|
||||
super().__init__()
|
||||
|
||||
# size: 回归 (log-transformed)
|
||||
size_dropout = min(0.5, dropout + 0.2)
|
||||
self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
|
||||
|
||||
# PDI: 2 分类
|
||||
self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# Encapsulation Efficiency: 3 分类
|
||||
self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# quantified_delivery: 回归 (z-scored)
|
||||
self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout)
|
||||
|
||||
# Biodistribution: 7 输出，softmax (sum=1)
|
||||
self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# toxic: 二分类
|
||||
self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0
|
||||
self.log_vars = nn.ParameterDict({
|
||||
t: nn.Parameter(torch.zeros(()))
|
||||
for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
|
||||
})
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
x: [B, in_dim] fusion 层输出
|
||||
|
||||
Returns:
|
||||
Dict with keys:
|
||||
- "size": [B, 1]
|
||||
- "pdi": [B, 2] logits
|
||||
- "ee": [B, 3] logits
|
||||
- "delivery": [B, 1]
|
||||
- "biodist": [B, 7] probabilities (sum=1)
|
||||
- "toxic": [B, 2] logits
|
||||
"""
|
||||
return {
|
||||
"size": self.size_head(x),
|
||||
"pdi": self.pdi_head(x),
|
||||
"ee": self.ee_head(x),
|
||||
"delivery": self.delivery_head(x),
|
||||
"biodist": self.biodist_head(x),
|
||||
"toxic": self.toxic_head(x),
|
||||
}
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class RegressionHead(nn.Module):
|
||||
"""回归任务头：输出单个 float 值"""
|
||||
|
||||
def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, 1]"""
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class ClassificationHead(nn.Module):
|
||||
"""分类任务头：输出 logits"""
|
||||
|
||||
def __init__(
|
||||
self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, num_classes),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, num_classes] (logits)"""
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class DistributionHead(nn.Module):
|
||||
"""分布任务头：输出和为 1 的概率分布（用于 Biodistribution）"""
|
||||
|
||||
def __init__(
|
||||
self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(in_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, num_outputs),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""[B, in_dim] -> [B, num_outputs] (softmax, sum=1)"""
|
||||
logits = self.net(x)
|
||||
return F.softmax(logits, dim=-1)
|
||||
|
||||
|
||||
class MultiTaskHead(nn.Module):
|
||||
"""
|
||||
多任务预测头，根据任务配置自动创建对应的子头。
|
||||
|
||||
输出:
|
||||
- size: [B, 1] 回归
|
||||
- pdi: [B, 2] 分类 logits
|
||||
- ee: [B, 3] 分类 logits
|
||||
- delivery: [B, 1] 回归
|
||||
- biodist: [B, 7] softmax 分布
|
||||
- toxic: [B, 2] 二分类 logits
|
||||
"""
|
||||
|
||||
def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
|
||||
super().__init__()
|
||||
|
||||
# size: 回归 (log-transformed)
|
||||
size_dropout = min(0.5, dropout + 0.2)
|
||||
self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
|
||||
|
||||
# PDI: 2 分类
|
||||
self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# Encapsulation Efficiency: 3 分类
|
||||
self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# quantified_delivery: 回归 (z-scored)
|
||||
self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout)
|
||||
|
||||
# Biodistribution: 7 输出，softmax (sum=1)
|
||||
self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# toxic: 二分类
|
||||
self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
|
||||
|
||||
# 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0
|
||||
self.log_vars = nn.ParameterDict({
|
||||
t: nn.Parameter(torch.zeros(()))
|
||||
for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
|
||||
})
|
||||
|
||||
def forward(self, x: torch.Tensor, x_numeric: torch.Tensor = None) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
x: [B, in_dim] fusion 层输出
|
||||
|
||||
Returns:
|
||||
Dict with keys:
|
||||
- "size": [B, 1]
|
||||
- "pdi": [B, 2] logits
|
||||
- "ee": [B, 3] logits
|
||||
- "delivery": [B, 1]
|
||||
- "biodist": [B, 7] probabilities (sum=1)
|
||||
- "toxic": [B, 2] logits
|
||||
"""
|
||||
xn = x if x_numeric is None else x_numeric
|
||||
return {
|
||||
"size": self.size_head(xn),
|
||||
"pdi": self.pdi_head(x),
|
||||
"ee": self.ee_head(x),
|
||||
"delivery": self.delivery_head(xn),
|
||||
"biodist": self.biodist_head(x),
|
||||
"toxic": self.toxic_head(x),
|
||||
}
|
||||
|
||||
@ -153,4 +153,7 @@ class ResidualConcatFusion(nn.Module):
|
||||
if f_retr is not None:
|
||||
out = out + self.g_retr * f_retr # 检索旁路，零初始化门保证起点=不开
|
||||
|
||||
return (out, attn) if return_attn_weights else out
|
||||
# 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr）
|
||||
if return_attn_weights:
|
||||
return out, pooled, attn
|
||||
return out, pooled
|
||||
@ -99,6 +99,8 @@ class LNPModel(nn.Module):
|
||||
moe_top_k: int = 2,
|
||||
moe_expert_hidden_mult: int = 2,
|
||||
moe_jitter_noise: float = 0.0,
|
||||
# ============ 回归旁路开关 ============
|
||||
reg_bypass: str = "on",
|
||||
# ============ LLM 相关 ============
|
||||
use_llm: bool = False,
|
||||
llm_model_path: str = DEFAULT_MOLT5_PATH,
|
||||
@ -174,6 +176,7 @@ class LNPModel(nn.Module):
|
||||
|
||||
# ============ LLM Prompt (可选) ============
|
||||
self.use_llm = use_llm
|
||||
self.reg_bypass = (str(reg_bypass).lower() == "on")
|
||||
if use_llm:
|
||||
self.llm_prompt = LLMPromptEncoder(
|
||||
d_model=d_model,
|
||||
@ -270,7 +273,9 @@ class LNPModel(nn.Module):
|
||||
_feats_t = torch.as_tensor(_feats, dtype=chem.dtype, device=chem.device)
|
||||
f_retr = self.retr_proj(_feats_t)
|
||||
|
||||
return self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
|
||||
fused, pooled = self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
|
||||
self._last_pooled = pooled # 纯数值向量，供回归 head 使用
|
||||
return fused
|
||||
|
||||
def forward_from_projected(
|
||||
self,
|
||||
@ -357,7 +362,7 @@ class LNPModel(nn.Module):
|
||||
delivery [B,1], biodist [B,7], toxic [B,2]
|
||||
"""
|
||||
fused = self.forward_backbone(smiles, tabular)
|
||||
return self.head(fused)
|
||||
return self.head(fused, getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None)
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""清空所有 encoder 的缓存"""
|
||||
@ -443,6 +448,8 @@ class LNPModelWithoutMPNN(LNPModel):
|
||||
moe_top_k: int = 2,
|
||||
moe_expert_hidden_mult: int = 2,
|
||||
moe_jitter_noise: float = 0.0,
|
||||
# ============ 回归旁路开关 ============
|
||||
reg_bypass: str = "on",
|
||||
# ============ LLM 相关 ============
|
||||
use_llm: bool = False,
|
||||
llm_model_path: str = DEFAULT_MOLT5_PATH,
|
||||
@ -472,6 +479,7 @@ class LNPModelWithoutMPNN(LNPModel):
|
||||
mpnn_checkpoint=None,
|
||||
mpnn_ensemble_paths=None,
|
||||
input_dims=dims,
|
||||
reg_bypass=reg_bypass,
|
||||
use_moe=use_moe,
|
||||
moe_n_experts=moe_n_experts,
|
||||
moe_top_k=moe_top_k,
|
||||
|
||||
@ -930,6 +930,8 @@ def main(
|
||||
moe_jitter_noise: float = 0.0,
|
||||
# Set Transformer
|
||||
set_transformer_block: str = "sab",
|
||||
# 回归旁路（消融开关）
|
||||
reg_bypass: str = "on",
|
||||
# LLM（消融开关）
|
||||
use_llm: bool = False,
|
||||
use_rag: bool = False,
|
||||
@ -967,6 +969,7 @@ def main(
|
||||
set_global_seed(seed)
|
||||
|
||||
llm_kwargs = dict(
|
||||
reg_bypass=reg_bypass,
|
||||
use_rag=use_rag,
|
||||
rag_top_k=rag_top_k,
|
||||
use_llm=use_llm,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user