简介一套面向计算机相关专业课程设计与期末大作业的Python项目基于知识图谱与推荐系统预测药物靶点交互曾获98分并经导师指导通过。项目完整整合生物医学知识图谱构建与推荐算法流程覆盖数据预处理、图谱建模、模型训练与结果评估等环节代码中提供hetionet、luo.s_dataset、yamanishi_08等经典数据集处理脚本以及deepdti、kge_rf、kge_nfm三套推荐模型实现可直接运行train_all.py复现预测实验。资源包共40个文件以9个.py源码脚本为核心辅以requirements.txt、Pipfile/Pipfile.lock依赖配置、README.md说明与操作指南压缩包仅56KB目录结构清晰。目前已有95人浏览学习适合需要快速上手知识图谱与推荐系统联合应用的开发者。1. 药物靶点交互预测为什么知识图谱加推荐系统能当成一个推荐问题来做做药物重定位的人最怕的事是在实验台上烧了三个月经费最后发现选中的靶点根本不表达。湿实验成本太高所以大家普遍先做计算筛选把成千上万个候选缩到几十个再进实验室。所谓药物靶点交互预测英文简称 DTI 预测就是把这个筛选过程数据化把已知的药物-蛋白关系、蛋白-疾病关系、药物-副作用关系全部织成一张图然后让模型回答「这个药和这个靶点之间会不会存在相互作用」。这张图本质上就是知识图谱而「会不会存在相互作用」在数学上又等价于推荐系统里的「用户会不会点这个商品」。知识图谱加推荐系统不是两个炫技名词的拼凑它们刚好各自解决了一半问题。下面我会把从数据准备、知识图谱构建到模型训练、候选预测的完整 Python 代码链路写清楚适合想快速复现一版能跑的 DTI 预测管线的人。代码不追求刷榜追求的是你能在本地把它改造成自己的数据能用的形状。先立住建模逻辑再动手写代码最后讲排错和验证。2. 建模逻辑把「药物-靶点交互」拆成推荐系统的链路预测问题2.1 三元组是知识图谱的最小单元DTI 预测本质是边补全知识图谱的最小单元是三元组(head, relation, tail)。放到药物这个域里DrugBank 上一条经过人工整理的条目可以写成(DB00001, has_target, P23219)表示这个药物分子和某个蛋白之间存在经过验证的相互作用。先把这句话拆开DB00001是药物实体P23219是靶点蛋白实体has_target是关系。图谱里所有信息都被压成这种「头实体-关系-尾实体」的扁平结构这也是它能把不同来源数据统一到一起的原因——不管来自 DrugBank、UniProt 还是 STRING最后都落成同样三列。现在换到推荐系统语境重新看这张表drug就是用户protein就是物品已知的has_target边就是用户给商品打的「正反馈」。目标从「推荐用户最可能点击的商品」变成「给这个药物推荐最可能结合的蛋白」。整个 DTI 预测任务本质上就是在这张异构图里做边补全已经有的边来自标注数据库没画出来的边就是我们需要打分的候选。所以我说知识图谱和推荐系统是各解决一半问题。知识图谱负责把多源数据组织成结构化拓扑推荐系统负责提供一套成熟的「给未知交互打分」的方法论。两者结合后的通用做法是把实体映射成 embedding然后用一个解码函数对(药物, 关系, 靶点)打分分数越高代表这条交互越可能存在。下面这张表是 DTI 图谱里最常见的实体和关系设计照着这个 schema 准备数据不会走偏。实体/关系头实体类型尾实体类型典型来源语义has_targetdrugproteinDrugBank / BindingDB药物直接靶向该蛋白treatsdrugdiseaseCTD该药用于治疗该疾病interacts_withproteinproteinSTRING蛋白间存在互作has_side_effectdrugside_effectSIDER该药具有某副作用2.2 药物就是用户靶点就是物品为什么内积建模在这个领域成立如果手上只有 drug-protein 交互矩阵推荐系统里最经典的 SVD 矩阵分解现在就能直接套用。把交互矩阵 M 近似成两个低秩矩阵的乘积 U·V^TU 每行是药物在隐空间的向量V 每列是蛋白的向量点积就是交互概率。这个建模假设在药理学里其实是成立的结构相似的化合物倾向于结合相似的靶点对应推荐里的「相似用户喜欢相似物品」。所以你不能说这是生搬硬套它背后有 Matched Molecular Pairs 这类药物化学里被反复验证的观察支撑。但纯矩阵分解有个明显短板药物和靶点的交互极其稀疏。一个冷门靶点如果只有两三条已知交互V 向量基本学不出有效信息一个没有任何已知交互的新药U 向量只能落在随机初始化附近。知识图谱把 disease、side_effect、protein-protein 交互加进来等于给矩阵分解补上了额外的侧信息让节点向量不再只依赖交互这一列。我一般拿到数据第一件事不是建模而是先统计这个交互矩阵到底有多稀疏稀疏度直接决定后面的模型选型。import pandas as pd # dtis.tsv 格式: drug_id, protein_id两列之间用 tab 分隔 dti pd.read_csv(dtis.tsv, sep\t, headerNone, names[drug, protein]) n_drugs dti[drug].nunique() n_prots dti[protein].nunique() n_known len(dti) # 稀疏度 1 - 已知交互数 / 全组合数 sparsity 1 - n_known / (n_drugs * n_prots) print(f药物 {n_drugs} 个靶点 {n_prots} 个已知交互 {n_known} 条) print(f矩阵稀疏度 {sparsity:.4%})nunique()拿到的是去重后的实体数量len(dti)是原始交互行数两者相除算出的是「已知边占全排列的比例」。一个典型的 DTI 数据集里这个值通常小于 1%稀疏度超过 99%。这时你就知道纯矩阵分解的泛化能力会很差必须引入知识图谱侧信息。代码跑完如果发现比值大于 5%说明你的数据集非常稠密那后面 TransE 未必是最优选择可以优先考虑直接上 LightGCN 这类图协同过滤。2.3 选型先 TransE再考虑图神经网络最后才上复杂编码器推荐的常见技术路线有三条矩阵分解、知识图谱嵌入、图神经网络。DTI 预测我建议从知识图谱嵌入里的 TransE 起步而不是一上来就上 RGCN 或 KGAT。原因有三。第一DTI 数据量通常只有几千到几万条三元组图神经网络需要多层消息传递在小图上非常容易过拟合。TransE 只有两个嵌入表和加法运算几十轮就能稳定收敛。第二TransE 的打分函数是||h r - t||给结果留了解释空间可以取h r与目标蛋白做最近邻直接回答「为什么这个药物会预测到蛋白 P」。第三RGCN 的优越性主要来自对不同关系类型做单独变换矩阵当你手里关系只有四五个、每种关系又只有零星几条时变换矩阵根本学不稳。只有当数据量真的到了上万条三元组并且has_target、inhibits、activates这些关系确实在语义上差异明显时才值得换 RGCN。提示选型不要被「最新模型」带着走。DTI 图谱的瓶颈通常是数据量和标注质量而不是模型容量。小数据集上 TransE 的效果未必比 RGCN 差而且翻车时更好排查。3. 数据准备从多源数据库构造知识图谱、邻接矩阵与负样本3.1 把四类原始表合成统一的三元组构建知识图谱的第一步是把不同来源的数据统一成三列。常见做法是去公开数据库下载四个表药物-靶点交互表、蛋白-蛋白互作表、药物-疾病关系表、药物-副作用表。下载之后你不要急着合并先做三件事统一 ID、去重、列名标准化。比如同一个蛋白在不同来源里可能叫P23219也可能叫PTGS1_HUMAN不统一后面映射全是洞。下面的代码把四个来源的表读进来各自打上 relation 标签然后重命名列后堆叠成一个三列的大表。注意pd.concat之前每个子表都只保留两列否则后面按列名取名会踩空。import pandas as pd dpi pd.read_csv(drug_protein.tsv, sep\t, names[drug, target]) ppi pd.read_csv(protein_protein.tsv, sep\t, names[p1, p2]) ddi pd.read_csv(drug_disease.tsv, sep\t, names[drug, disease]) sider pd.read_csv(drug_side_effect.tsv, sep\t, names[drug, se]) # 统一加上关系类型 dpi[relation] has_target ppi[relation] interacts_with ddi[relation] treats sider[relation] has_side_effect triplets pd.concat([ dpi.rename(columns{drug: head, target: tail}), ppi.rename(columns{p1: head, p2: tail}), ddi.rename(columns{drug: head, disease: tail}), sider.rename(columns{drug: head, se: tail}), ], ignore_indexTrue)[[head, relation, tail]] # 去重防止同一来源重复条目污染统计 triplets triplets.drop_duplicates() print(triplets.shape)rename把每个子表的列统一成head和tailrelation是提前打好的关系标签。drop_duplicates()这一步别省不同数据库之间有大量冗余关系同一对实体出现三四次很常见。打印shape时注意行数变化如果合并后行数比原始表加起来还多一定是列名没对齐导致出现了笛卡尔积。3.2 实体映射与邻接矩阵用 COO 格式而不是稠密矩阵计算机不认字符串训练前要把所有实体和关系映射成整数 ID然后构造图的边索引。这里我要强调一个新手高频翻车点如果你用现有 Python 包里的稠密邻接矩阵实体数一到几千就会内存爆炸。几千实体对应的稠密矩阵是几千万个格子纯浮点数直接上百 MB模型还没训就先把内存吃光了。正确做法是只记录有边的位置也就是 COO 稀疏格式。# 实体与关系统一编 ID entities sorted(set(triplets[head]) | set(triplets[tail])) ent2id {e: i for i, e in enumerate(entities)} rel2id {r: i for i, r in enumerate(sorted(triplets[relation].unique()))} triplets[head_id] triplets[head].map(ent2id) triplets[tail_id] triplets[tail].map(ent2id) triplets[rel_id] triplets[relation].map(rel2id) # 建边正反向都写入保证图传播时信息能双向流动 edges pd.concat([ triplets[[head_id, tail_id, rel_id]], triplets[[tail_id, head_id, rel_id]], # 反向边 ], ignore_indexTrue) n_ent len(ent2id) n_rel len(rel2id) print(f实体数 {n_ent}关系数 {n_rel}边数 {len(edges)})反向边是知识图谱建模里很关键的一步。has_target关系如果只保留药物指向蛋白的方向图神经网络的消息传递就只能从药物流到蛋白蛋白侧的邻居信息传不回来。我们做图谱构建时把每条边正反都写一份相当于把有向图变成无向图来传播特征。对 DTI 预测来说相互作用本身无方向这个处理同时符合生物语义和工程需要。3.3 负采样策略随机替换头实体或尾实体训练集不能只有正样本模型必须有「这条边不存在」的例子才能学会区分。推荐系统里这叫负采样知识图谱里习惯叫随机破坏三元组。常见做法是等概率地把头实体或尾实体随机替换成另一个实体替换后要检查新生成的三元组没有出现在正样本集里否则就成了假负样本。import random def negative_sample(triplets, ent2id, neg_ratio2, seed42): rng random.Random(seed) heads triplets[head_id].to_numpy() tails triplets[tail_id].to_numpy() rels triplets[rel_id].to_numpy() n_ent len(ent2id) known set(zip(heads, rels, tails)) # 正样本集用于过滤 neg [] for h, r, t in zip(heads, rels, tails): for _ in range(neg_ratio): while True: if rng.random() 0.5: nh rng.randrange(n_ent) cand (nh, r, t) else: nt rng.randrange(n_ent) cand (h, r, nt) if cand not in known: break neg.append(cand) return negneg_ratio2表示每条正样本生成 2 条负样本常见取值在 1 到 3 之间太高会让模型偏向把所有边都判为不存在。known用集合而不是列表因为in判断在集合上是 O(1)列表一长会慢到怀疑人生。while True内层循环是为了避免随机替换恰好撞上正样本虽然概率低但数据集大时一定会发生不过滤就会在训练时给模型喂自相矛盾的标签。提示切分的正确顺序是先按药物或靶点把交互对切好再对训练子集做负采样。如果先对全量正样本做负采样再切分测试集正样本的负样本可能混进训练集离线指标会虚高得离谱这个坑在第 5 章专门讲。4. 用纯 PyTorch 跑通 TransE 模型从嵌入矩阵到候选预测4.1 定义 TransE 模型与打分函数TransE 的核心假设就一句话如果三元组成立头实体向量加上关系向量应该约等于尾实体向量。所以打分函数取h r - t的范数距离越小说明这条边越可能为真。为了让输出变成适合二分类的 logit我用一个边界值gamma减去距离距离为 0 时分数最高等于gamma距离越大分数越低。import torch import torch.nn as nn class TransE(nn.Module): def __init__(self, n_ent, n_rel, dim128, gamma6.0): super().__init__() self.ent_emb nn.Embedding(n_ent, dim) self.rel_emb nn.Embedding(n_rel, dim) self.gamma gamma def forward(self, heads, rels, tails): h self.ent_emb(heads) # [batch, dim] r self.rel_emb(rels) # [batch, dim] t self.ent_emb(tails) # [batch, dim] # 距离用 L2 范数小数据集上比 L1 更稳 d (h r - t).norm(p2, dim-1) return self.gamma - d嵌入维度dim128是个弹性参数DTI 小数据集上 64 到 128 都常见。维度太低表达不了实体语义太高在小数据上容易过拟合。gamma是距离向 logit 映射的缩放边界我一般初始化 6.0如果 loss 不下降先看gamma和学习率而不是盲目加模型深度。注意这里没有 sigmoid因为后面搭配的BCEWithLogitsLoss内部自带 sigmoid不要再手动加激活函数。4.2 组装训练数据与训练循环有了模型下一步把正负三元组组装成训练张量。正样本标签为 1负样本标签为 0然后包成TensorDataset给DataLoader分批。训练循环就是标准的 PyTorch 套路前向打分、算 loss、反向传播、更新参数。from torch.utils.data import DataLoader, TensorDataset pos torch.tensor( triplets[[head_id, rel_id, tail_id]].values, dtypetorch.long ) neg torch.tensor( negative_sample(triplets, ent2id, neg_ratio2), dtypetorch.long ) all_tri torch.cat([pos, neg], dim0) labels torch.cat([torch.ones(len(pos)), torch.zeros(len(neg))], dim0) ds TensorDataset(all_tri[:, 0], all_tri[:, 1], all_tri[:, 2], labels) loader DataLoader(ds, batch_size256, shuffleTrue) model TransE(n_entn_ent, n_reln_rel, dim128) opt torch.optim.Adam(model.parameters(), lr1e-3) loss_fn nn.BCEWithLogitsLoss() for epoch in range(50): model.train() total_loss 0.0 for h, r, t, y in loader: logits model(h, r, t) loss loss_fn(logits, y) opt.zero_grad() loss.backward() opt.step() total_loss loss.item() print(fepoch {epoch:02d}, loss {total_loss:.4f})negative_sample返回元组列表直接转torch.long张量没问题。batch_size256是常用起点显存紧张就降到 128。lr1e-3配 Adam 在嵌入模型上表现稳定50 轮足够小数据集收敛。如果你想试试正负样本比例对结果的影响把neg_ratio从 2 改成 3 或 5观察 loss 的收敛平台和评估指标变化——这是最便宜的消融实验。4.3 推理阶段对全部候选药物-靶点对打分训练结束后把测试集里所有尚未标注的「药物 × 靶点」组合拿出来逐批喂给模型打分然后按分数排序输出候选表。这里有三个细节第一推理时relation固定为has_target的 ID因为我们只关心 DTI第二必须分批计算全量组合可能上百万条一次性前向会撑爆显存第三打分前model.eval()并包在torch.no_grad()里省内存也避免 dropout 造成干扰。drug_ids torch.tensor(list(drug_set), dtypetorch.long) protein_ids torch.tensor(list(protein_set), dtypetorch.long) # 生成全部候选对 candidates torch.cartesian_prod(drug_ids, protein_ids) has_target_rel rel2id[has_target] scores [] bs 4096 with torch.no_grad(): model.eval() for start in range(0, len(candidates), bs): h candidates[start:start bs, 0] t candidates[start:start bs, 1] r torch.full_like(h, has_target_rel) scores.append(model(h, r, t).cpu()) scores torch.cat(scores) res pd.DataFrame({ drug_id: candidates[:, 0].numpy(), protein_id: candidates[:, 1].numpy(), score: scores.numpy(), }).sort_values(score, ascendingFalse) res.head(100).to_csv(dti_top100_candidates.tsv, sep\t, indexFalse)torch.full_like(h, has_target_rel)给每个候选对都填上同一个关系 ID这样模型只会在 DTI 这个关系上打分。sort_values按分数从高到低排序head(100)直接输出前 100 个候选。注意这个候选列表里混着训练集见过的边正式评估时要先剔除否则 Top 结果全是已知交互的复现没有预测价值。5. 参数与踩坑DTI 预测最常见的 5 个翻车点5.1 邻接矩阵变稠密内存直接爆掉现象代码跑到构建图那一步进程被系统杀掉或者MemoryError。原因新手图省事用np.zeros((n_ent, n_ent))建稠密矩阵。实体数 5000 时就是 2500 万个格子每个浮点数 8 字节一次就是 200 MB。加上后续还要存多张关系矩阵内存翻几倍就崩了。解决全部改成稀疏 COO 格式。PyTorch 里可以用torch.sparse_coo_tensor只存有边位置的行列索引和值。row torch.tensor(edges[head_id].values, dtypetorch.long) col torch.tensor(edges[tail_id].values, dtypetorch.long) value torch.ones(len(row)) adj torch.sparse_coo_tensor(torch.stack([row, col]), value, (n_ent, n_ent))注意矩阵维度是(n_ent, n_ent)不是(n_drugs, n_proteins)。很多人把「药物-靶点交互矩阵」和「全图谱邻接矩阵」搞混前者只有两类节点且是稠密的后者必须覆盖所有实体类型且是稀疏的。5.2 负采样泄漏AUC 虚高到 0.99现象测试集 AUC 高达 0.98、0.99看起来效果惊人但人工抽看 Top 候选全是已知交互的重复或同义改写。原因负采样在数据切分之前做了。测试集里的某个药物-靶点正样本在构造训练负样本时被随机替换成了一条一模一样的边训练时模型见过测试时当然认得。解决严格按「先切分、后负采样」的顺序。先把正样本按药物或靶点切开保证测试交互对不出现在训练三元组里然后只对训练部分做负采样并且负样本生成后还要过滤一遍把落在测试集正样本集合里的负样本删掉。train_pairs set(zip(train[drug_id], train[protein_id])) test_pairs set(zip(test[drug_id], test[protein_id])) assert not (train_pairs test_pairs), 训练集与测试集存在交互对泄漏assert这行是白送的防线。每次切完数据先跑一遍两个集合交集不为空就立刻停。血泪经验是这个检查能拦住 80% 的 DTI 评估失真比后面任何调参都值钱。5.3 实体映射出现 NaN模型索引越界现象map之后打印head_id出现nan训练时 Embedding 查表报 index 越界或者 loss 直接变 NaN。原因某条数据的实体 ID 在字典里不存在map默认返回 NaN。NaN 被当成 float 塞进torch.long张量时要么转成垃圾整数要么直接报错。最常见于不同来源数据库之间 ID 命名不一致。解决建图前先做完整性检查。映射后马上断言没有缺失值缺失的实体要么补字典要么删行不能留到训练阶段。import math assert triplets[head_id].notna().all(), 存在未映射的头实体 assert triplets[tail_id].notna().all(), 存在未映射的尾实体 assert not any(math.isnan(x) for x in triplets[rel_id]), 存在未映射的关系如果断言失败排查方向是去原始表里看那个实体的原始 ID 长什么样。常见情况是某个表里用了版本化 ID 后缀另一个表里没有去掉后缀再映射就通了。5.4 新药冷启动预测结果全是噪声现象某个没有任何已知交互的新药模型给出的 Top 靶点杂乱无章跟药物化学直觉对不上。原因TransE 这类嵌入方法节点向量完全靠图中的边学习。新加入的实体一条边都没有嵌入始终停留在随机初始化附近打分当然没有意义。解决给孤立节点补上特征边。常见做法有两条一是把化学结构相似性转换成语义边最相似的药物之间加一条similar_to关系二是用分子指纹或蛋白序列特征做节点初始化让嵌入表至少有先验信息。完全没有任何边的实体我一般直接从候选列表剔除不做无依据预测在给业务方的说明里更诚实。5.5 只盯着 AUCPR-AUC 低到没法看现象AUC 0.90 以上但把 Top 50 候选拿去湿实验验证命中率不到 10%。原因DTI 数据天然正负样本极不平衡负样本占绝大多数。AUC 对类别不平衡不敏感它度量的是「随机正样本排在随机负样本前面的概率」哪怕排序质量很差只要负样本够多也能维持高值。解决评估指标换成 PR-AUC、RecallK、HitK。PR-AUC 对正样本的精确率和召回率更敏感K 的取值要和你的湿实验预算对齐——预算只够验证 30 个就看 Recall30。from sklearn.metrics import average_precision_score # test_scores: 模型对测试集输出的分数 # test_labels: 测试集真实标签 pr_auc average_precision_score(test_labels, test_scores) print(fPR-AUC {pr_auc:.4f})average_precision_score就是 PR-AUC把预测分数当输入不需要先设阈值。如果这个值明显低于 AUC说明模型对正样本的排序不够靠前得回去调负采样比例或嵌入维度而不是换一个更复杂的模型。6. 从候选列表到可验证假设AUC 之外的三步验证方法离线指标再漂亮也只是给湿实验做漏斗过滤。我第一次跑通这套流程时指标相当好看兴冲冲把 Top 候选发给做药理的同学结果被反问一句「这里有几个已经被别的实验室验证过了」当场卡住。后来我养成一个习惯每次预测结果出来先做三步人工验证再往外发。第一步是相似物证据链。把候选药物在嵌入空间里的最近邻药物找出来看它是否有已知的靶点证据。TransE 的向量空间本身保留了「相似实体距离近」的性质候选药物最近的邻居如果是某个已有明确 DTI 标注的已知药物这份证据可以直接写进报告。第二步是图谱路径反查。用 networkx 或直接遍历邻接表找出候选药物到目标蛋白之间是否存在一条通顺的三元组路径比如「药物 A → treats → 疾病 D → 该靶点参与疾病通路」。路径越短、中间节点越少这个假设越适合优先验证。第三步是拿外部独立来源交叉核对把 Top 候选里的蛋白 ID 拿到其他线上数据库反查最近有没有新的实验证据如果外部已有报道至少说明这不是纯计算噪声。这套三步验证法看起来不 fancy但它才是把模型输出转成实验决策的关键环节。模型负责把一万个候选缩到一百个人工路径验证负责把一百个缩到十个。每次跑完预测我把 Top 100 连同最近邻、最短路径、外部证据状态一起存成一个verification_report.tsv评审会直接看表就行。希望这套从数据到模型再到验证的路径能帮你在自己的项目上少走弯路——毕竟计算预测的价值永远要落在实验台上才作数。本文还有配套的精品资源点击获取
