import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm import torch.nn.functional as F import torch.distributed as dist from model import load_encoder_components, ProteinMoleculeDualEncoder from dataset import ProteinMoleculeDataset, DualTowerCollator def print_model_stats(model): total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) # 格式化数字，例如转换为 Million (M) def format_params(num): if num >= 1e6: return f"{num / 1e6:.2f}M" elif num >= 1e3: return f"{num / 1e3:.2f}K" else: return str(num) print(f"|" + "-"*50 + "|") print(f"| {'Model Statistics':^48} |") print(f"|" + "-"*50 + "|") print(f"| Total Parameters : {format_params(total_params):>15} |") print(f"| Trainable Parameters : {format_params(trainable_params):>15} |") print(f"| Frozen Parameters : {format_params(total_params - trainable_params):>15} |") print(f"|" + "-"*50 + "|") # 分别查看两个塔的大小（可选，帮你确认 backbone 大小） if hasattr(model, 'protein_encoder'): p_params = sum(p.numel() for p in model.protein_encoder.parameters()) print(f"| Protein Tower : {format_params(p_params):>15} |") if hasattr(model, 'molecule_encoder'): m_params = sum(p.numel() for p in model.molecule_encoder.parameters()) print(f"| Molecule Tower : {format_params(m_params):>15} |") print(f"|" + "-"*50 + "|") class DualTowerTrainer: def __init__( self, model: nn.Module, train_loader: DataLoader, val_loader: DataLoader, learning_rate: float = 1e-4, temperature: float = 0.07, device: str = 'cuda' if torch.cuda.is_available() else 'cpu', save_dir: str = "./checkpoints", ): print_model_stats(model) self.model = model.to(device) self.train_loader = train_loader self.val_loader = val_loader self.device = device self.temperature = temperature self.save_dir = save_dir self.optimizer = optim.AdamW(self.model.parameters(), lr=learning_rate) self.scheduler = optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_max=10, eta_min=1e-6 ) if not os.path.exists(save_dir): os.makedirs(save_dir) def compute_loss(self, p_vec, m_vec, labels): """ 计算 Masked InfoNCE Loss。 只对 batch 中 label=1 (Active) 的样本对计算正向损失。 所有样本 (包括 label=0) 都会作为分母中的负例。 """ # 1. 计算相似度矩阵 (Batch_Size, Batch_Size) # logits[i][j] 表示第 i 个蛋白和第 j 个分子的相似度 logits = torch.matmul(p_vec, m_vec.T) / self.temperature # 2. 生成目标 (对角线是正例) batch_size = p_vec.size(0) targets = torch.arange(batch_size).to(self.device) # 3. 计算 Cross Entropy # 我们只关心 label=1 的行，因为 label=0 的行本身就不应该结合， # 如果强制 label=0 的对角线相似度最大化是错误的。 # CrossEntropyLoss 默认是会对 logits 进行 Softmax # 这里的 loss 是计算每一行（每个 Protein）去匹配正确的 Molecule raw_loss = F.cross_entropy(logits, targets, reduction='none') # 4. Masking: 只取 active (label=1) 的 loss 平均 active_mask = (labels == 1).float() # 防止除以 0 num_actives = active_mask.sum() if num_actives > 0: final_loss = (raw_loss * active_mask).sum() / num_actives else: # 如果这一个 batch 全是负例，loss 为 0 (或者为了梯度流不断，设为一个很小的值) # final_loss = torch.tensor(0.0, device=self.device, requires_grad=True) final_loss = 0.0 * (p_vec.sum() + m_vec.sum()) return final_loss def train_epoch(self, epoch_idx): self.model.train() total_loss = 0 active_count = 0 loop = tqdm(self.train_loader, desc=f"Train Epoch {epoch_idx}") for batch in loop: # 1. 数据移到 GPU prot_inputs = {k: v.to(self.device) for k, v in batch['protein_inputs'].items()} mol_inputs = {k: v.to(self.device) for k, v in batch['molecule_inputs'].items()} labels = batch['labels'].to(self.device) # 0 或 1 if (labels == 1).sum() == 0: print("\nSkipping batch with no active samples") continue # 2. 前向传播 self.optimizer.zero_grad() p_vec, m_vec = self.model(prot_inputs, mol_inputs) # 3. 计算 Loss loss = self.compute_loss(p_vec, m_vec, labels) # 4. 反向传播 loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) # 梯度裁剪防止爆炸 self.optimizer.step() total_loss += loss.item() if (labels == 1).sum() > 0: active_count += 1 loop.set_postfix(loss=loss.item()) avg_loss = total_loss / len(self.train_loader) self.scheduler.step() # 更新学习率 return avg_loss @torch.no_grad() def evaluate(self): self.model.eval() total_loss = 0 # 也可以统计准确率：对于 Active 的对子，Top-1 预测是不是它自己 correct_retrieval = 0 total_actives = 0 for batch in tqdm(self.val_loader, desc="Evaluating"): prot_inputs = {k: v.to(self.device) for k, v in batch['protein_inputs'].items()} mol_inputs = {k: v.to(self.device) for k, v in batch['molecule_inputs'].items()} labels = batch['labels'].to(self.device) p_vec, m_vec = self.model(prot_inputs, mol_inputs) loss = self.compute_loss(p_vec, m_vec, labels) total_loss += loss.item() # 计算简单的 Top-1 Accuracy (仅针对 Active 样本) logits = torch.matmul(p_vec, m_vec.T) preds = torch.argmax(logits, dim=1) # 每一行预测最大的列索引 targets = torch.arange(p_vec.size(0)).to(self.device) active_mask = (labels == 1) if active_mask.sum() > 0: matches = (preds[active_mask] == targets[active_mask]) correct_retrieval += matches.sum().item() total_actives += active_mask.sum().item() avg_loss = total_loss / len(self.val_loader) acc = correct_retrieval / total_actives if total_actives > 0 else 0.0 return avg_loss, acc # def save_checkpoint(self, epoch, metric): # path = os.path.join(self.save_dir, f"model_epoch_{epoch}_acc_{metric:.4f}.pt") # torch.save(self.model.state_dict(), path) # print(f"Model saved to {path}") def save_checkpoint(self, epoch, metric): path = os.path.join(self.save_dir, f"model_epoch_{epoch}_acc_{metric:.4f}.pt") # 检查模型是否被 DDP 或 DataParallel 包裹 if isinstance(self.model, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)): state_dict = self.model.module.state_dict() else: state_dict = self.model.state_dict() torch.save(state_dict, path) if dist.get_rank() == 0: print(f"Model saved to {path} (Clean state_dict without 'module.' prefix)") def train_model( dataset_path: str, model_and_tokenizers = None, protein_model_path: str = None, molecule_model_path: str = None, model_save_dir: str = "./output_checkpoints", epochs: int = 10, batch_size: int = 32, lr: float = 1e-4 ): # 1. 加载 Tokenizer 和 Dataset print("Initialize components...") if model_and_tokenizers is not None: model, p_tokenizer, m_tokenizer = model_and_tokenizers else: p_encoder, p_tokenizer, m_encoder, m_tokenizer = load_encoder_components( protein_model_path, molecule_model_path ) model = ProteinMoleculeDualEncoder(p_encoder, m_encoder, projection_dim=256) full_dataset = ProteinMoleculeDataset(dataset_path) # 划分训练集和验证集 (80/20) train_size = int(0.8 * len(full_dataset)) val_size = len(full_dataset) - train_size train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size]) # 3. 准备 Collator 和 DataLoader collator = DualTowerCollator(p_tokenizer, m_tokenizer) train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collator, num_workers=4, # 根据CPU核心数调整 pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collator, num_workers=4 ) # 4. 初始化模型 # model = ProteinMoleculeDualEncoder(p_encoder, m_encoder, projection_dim=256) # 5. 初始化 Trainer trainer = DualTowerTrainer( model=model, train_loader=train_loader, val_loader=val_loader, learning_rate=lr, save_dir=model_save_dir, ) # 6. 开始循环 print("Start Training...") best_acc = 0.0 for epoch in range(epochs): train_loss = trainer.train_epoch(epoch) val_loss, val_acc = trainer.evaluate() print(f"Epoch {epoch} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Retrieval Acc: {val_acc:.4f}") # 保存最好模型 if val_acc > best_acc: best_acc = val_acc trainer.save_checkpoint(epoch, val_acc) if __name__ == "__main__": protein_model_path = "./SaProt_650M_AF2" molecule_model_path = "./ChemBERTa-zinc-base-v1" dataset_path = 'drug_target_activity/processed_train.parquet' model_save_dir = './Dual_Tower_Model/output_checkpoints' train_model( protein_model_path=protein_model_path, molecule_model_path=molecule_model_path, dataset_path=dataset_path, model_save_dir=model_save_dir )