☰
知识图谱与推荐系统融合:DTI药物靶点交互预测的TransE实现全解析
2026/10/26 16:02:11 网站建设 项目流程

简介:一套面向计算机相关专业课程设计与期末大作业的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^T,U 每行是药物在隐空间的向量,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", header=None, 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_index=True)[["head", "relation", "tail"]] # 去重,防止同一来源重复条目污染统计 triplets = triplets.drop_duplicates() print(triplets.shape)

rename把每个子表的列统一成head和tail,relation是提前打好的关系标签。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_index=True) 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_ratio=2, seed=42): 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 neg

neg_ratio=2表示每条正样本生成 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, dim=128, gamma=6.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(p=2, dim=-1) return self.gamma - d

嵌入维度dim=128是个弹性参数,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, dtype=torch.long ) neg = torch.tensor( negative_sample(triplets, ent2id, neg_ratio=2), dtype=torch.long ) all_tri = torch.cat([pos, neg], dim=0) labels = torch.cat([torch.ones(len(pos)), torch.zeros(len(neg))], dim=0) ds = TensorDataset(all_tri[:, 0], all_tri[:, 1], all_tri[:, 2], labels) loader = DataLoader(ds, batch_size=256, shuffle=True) model = TransE(n_ent=n_ent, n_rel=n_rel, dim=128) opt = torch.optim.Adam(model.parameters(), lr=1e-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(f"epoch {epoch:02d}, loss {total_loss:.4f}")

negative_sample返回元组列表,直接转torch.long张量没问题。batch_size=256是常用起点,显存紧张就降到 128。lr=1e-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), dtype=torch.long) protein_ids = torch.tensor(list(protein_set), dtype=torch.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", ascending=False) res.head(100).to_csv("dti_top100_candidates.tsv", sep="\t", index=False)

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, dtype=torch.long) col = torch.tensor(edges["tail_id"].values, dtype=torch.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 只盯着 AUC,PR-AUC 低到没法看

现象:AUC 0.90 以上,但把 Top 50 候选拿去湿实验验证,命中率不到 10%。

原因:DTI 数据天然正负样本极不平衡,负样本占绝大多数。AUC 对类别不平衡不敏感,它度量的是「随机正样本排在随机负样本前面的概率」,哪怕排序质量很差,只要负样本够多也能维持高值。

解决:评估指标换成 PR-AUC、Recall@K、Hit@K。PR-AUC 对正样本的精确率和召回率更敏感,K 的取值要和你的湿实验预算对齐——预算只够验证 30 个,就看 Recall@30。

from sklearn.metrics import average_precision_score # test_scores: 模型对测试集输出的分数 # test_labels: 测试集真实标签 pr_auc = average_precision_score(test_labels, test_scores) print(f"PR-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,评审会直接看表就行。希望这套从数据到模型再到验证的路径能帮你在自己的项目上少走弯路——毕竟计算预测的价值,永远要落在实验台上才作数。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询