首发于量子位
亚马逊上海AI实验室构建药物重定位「知识图谱」,帮助缩短新冠药物研发周期

亚马逊上海AI实验室构建药物重定位「知识图谱」,帮助缩短新冠药物研发周期

自2019年12月至今,新型冠状病毒在全球迅速扩散已超过720万例

目前急需快速有效的新冠病毒有效药物的发现路径。

药物重定位是一种将现有药物用于治疗新的适应症的药物发现方式。相比较传统的新药开发,它可以有效缩短药物研发周期,降低成本,规避风险。因此药物重定位是一种非常有前景的新冠肺炎治疗策略。

最近,亚马逊上海AI实验室联合来自亚马逊AI北美、明尼苏达大学、俄亥俄州立大学以及湖南大学的专家学者,共同构建了大规模「药物重定位知识图谱」**DRKG(**Drug Repurposing Knowledge Graph)以及一套完整的用于进行药物重定位研究的机器学习工具,并将其在github上开源。

什么是DRKG

DRKG是一个综合型生物医药知识图谱,其包含人类基因、化合物、生物过程、药物副作用、疾病和症状等六个主要方面的数据。

DRKG从六个公开的大型医药数据库以及近期新冠病毒的相关医学文献中挖掘并数据,并进行整理和规范。

DRKG知识图谱包含分属于13种实体类型97238个实体,以及分属于107种关系类型5874261个三元组数据。

Drug Repurposing Knowledge Graph

机器学习工具使用了先进的深度图学习方法(DGL-KE)来学习DRKG中实体和关系的低维向量表示(embeddings), 并使用这些embedding来预测药物治疗疾病的可能性或药物与疾病靶点结合的可能性。

DGL-KE是亚马逊上海AI实验室开源的一款专门针对大规模知识图谱嵌入表示的训练工具,同时支持单机多GPU和多机分布式训练。在AWS EC2平台上,一台p3.16xlarge(8 GPUs)可以在100分钟内训练完成Freebase数据集(8600万节点,3亿条边)。4台r5dn.24xlarge(4*48 CPUs) 可以在30分钟内完成训练。

DRKG在其开源代码仓库中提供了使用DGL-KE来生成DRKG知识图谱中实体和关系的低纬嵌入向量表示的案例以及基于预训练的DRKG知识图谱嵌入向量表示预测新冠病毒候选药物的案例。

初步的实验结果表明:使用此类机器学习工具进行冠状病毒的药物发现,能以较高的分数识别出目前正在进行临床试验的多种冠状病毒候选药物。

使用DGL-KE生成DRKG知识图谱的低纬嵌入向量表示

DGL-KE是一个基于DGL图深度学习框架,面向知识图谱嵌入学习方法领域的高性能、高可扩展的开源软件库,DRKG利用DGL-KE 软件包进行DRKG知识图谱的低纬嵌入向量表示的生成。

  • 下载DRKG知识图谱,DRKG知识图谱已开源,链接已附文末。


import sys
sys.path.insert(1, '../utils')
from utils import download_and_extract
download_and_extract()
drkg_file = '../data/drkg/drkg.tsv'
  • DRKG知识图谱包含一个tsv格式文件drkg.tsv,其中包含了知识图谱的所有三元组,在训练之前,我们将数据集随机按照0.9:0.05:0.05的比例划分成训练集、验证集和测试集。
import pandas as pd
import numpy as np
 
df = pd.read_csv(drkg_file, sep="\t")
triples = df.values.tolist()
seed = np.arange(num_triples)
np.random.shuffle(seed)
 
train_cnt = int(num_triples * 0.9)
valid_cnt = int(num_triples * 0.05)
train_set = seed[:train_cnt]
train_set = train_set.tolist()
valid_set = seed[train_cnt:train_cnt+valid_cnt].tolist()
test_set = seed[train_cnt+valid_cnt:].tolist()
 
with open("train/drkg_train.tsv", 'w+') as f:
   for idx in train_set:
       f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))
 
with open("train/drkg_valid.tsv", 'w+') as f:
   for idx in valid_set:
       f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))
 
with open("train/drkg_test.tsv", 'w+') as f:
   for idx in test_set:
       f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))
  • 随后直接调用DGL-KE软件包的命令行进行DRKG知识图谱的低纬嵌入向量表示训练,案例中我们选用TransE_l2知识图谱嵌入算法,并使用了AWS p3.16xlarge实例进行多GPU并行进行训练。
!DGLBACKEND=pytorch dglke_train --dataset DRKG --data_path ./train --data_files drkg_train.tsv drkg_valid.tsv drkg_test.tsv --format 'raw_udd_hrt' --model_name TransE_l2 --batch_size 2048 \--neg_sample_size 256 --hidden_dim 400 --gamma 12.0 --lr 0.1 --max_step 100000 --log_interval 1000 --batch_size_eval 16 -adv --regularization_coef 1.00E-07 --test --num_thread 1 --gpu 0 1 2 3 4 5 6 7 --num_proc 8 --neg_sample_size_eval 10000 --async_update
  • 训练完成后我们将得到两个文件: 1) DRKG_TransE_l2_entity.npy, DRKG中实体的低维向量表示和 2)DRKG_TransE_l2_relation.npy,DRKG中关系的低维向量表示。后续我们可以使用训练好的实体和关系的低维向量表示进行药物预测。
node_emb = np.load('./ckpts/TransE_l2_DRKG_0/DRKG_TransE_l2_entity.npy')
relation_emb = np.load('./ckpts/TransE_l2_DRKG_0/DRKG_TransE_l2_relation.npy')

使用预训练的DRKG知识图谱的嵌入向量预测新冠病毒候选药物

亚马逊AI联合湖南大学、克里夫兰诊所勒纳中心基因组医学研究所以及明尼苏达大学的专家学者的论文「Repurpose Open Data to Discover Therapeutics for COVID-19 using Deep Learning」提出了结合使用知识图谱的嵌入和基因集富集分析的方法来进行新冠病毒老药新用药物的预测。

DRKG基于该论文的思路,提供了基于预训练的DRKG知识图谱嵌入向量表示进行新冠病毒候选药物预测案例。

首先我们将基于DRKG知识图谱的新冠病毒的药物预测问题定义为预测药物和新冠病毒实体之间存在'Hetionet::CtD::Compound:Disease'‘GNBR::T::Compound:Disease'关系(即治疗treatment关系)的置信度评估问题。

我们选取DRKG中来自Drugbank的分子量(molecule weight)大于250的FDA获准药物实体作为可选药物集,同时选取DRKG知识图谱中存在的34个新冠病毒相关实体作为目标病毒集。

然后,我们预测所有可能的(药物,治疗,病毒)三元组组合在TrainsE_l2算法下的分数(score),并最终对分数进行排序,并选取分数最高的100个药物作为推荐药物。

具体方法实施如下:

  • 设定目标病毒实体、药物实体和治疗关系。
# 目标新冠病毒相关实体
COV_disease_list = ['Disease::SARS-CoV2 E','Disease::SARS-CoV2 M', ...]
 
# 药物疾病治疗相关关系
treatment = ['Hetionet::CtD::Compound:Disease','GNBR::T::Compound:Disease']
 
# 获取来自Drugbank的分子量(molecule weight)大于250的FDA获准药物实体(已在infer_drug.tsv中提供
drug_list = []
with open("./infer_drug.tsv", newline='', encoding='utf-8') as csvfile:
   reader = csv.DictReader(csvfile, delimiter='\t', fieldnames=['drug','ids'])
   for row_val in reader:
       drug_list.append(row_val['drug'])
  • 获取预训练DRKG知识图谱的嵌入表示
# 读取预训练embedding
entity_emb = np.load('../data/drkg/embed/DRKG_TransE_l2_entity.npy')
rel_emb = np.load('../data/drkg/embed/DRKG_TransE_l2_relation.npy')
 
drug_ids = th.tensor(drug_ids).long()
disease_ids = th.tensor(disease_ids).long()
treatment_rid = th.tensor(treatment_rid)
 
drug_emb = th.tensor(entity_emb[drug_ids])
treatment_embs = [th.tensor(rel_emb[rid]) for rid in treatment_rid]
  • 所有可能的(药物,治疗,病毒)三元组组合在TrainsE_l2算法下的分数(score),计算公式如下:

( h 为head 即药物; r 为关系; t 为tail即病毒;gamma为训练中使用的常数)

import torch.nn.functional as fn
 
gamma=12.0
def transE_l2(head, rel, tail):
   score = head + rel - tail
   return gamma - th.norm(score, p=2, dim=-1)
 
scores_per_disease = []
dids = []
# 针对两种治疗关系分别计算(药物,治疗,病毒)三元组的分数,并最终合并
for rid in range(len(treatment_embs)):
   treatment_emb=treatment_embs[rid]
   for disease_id in disease_ids:
       disease_emb = entity_emb[disease_id]
       score = fn.logsigmoid(transE_l2(drug_emb, treatment_emb, disease_emb))
       scores_per_disease.append(score)
       dids.append(drug_ids)
scores = th.cat(scores_per_disease)
 
  • 对分数进行排序
idx = th.flip(th.argsort(scores), dims=[0])
scores = scores[idx].numpy()
dids = dids[idx].numpy()
  • 获取最终topk的药物推荐
topk=100
_, unique_indices = np.unique(dids, return_index=True)
topk_indices = np.sort(unique_indices)[:topk]
# top100的药物ID
proposed_dids = dids[topk_indices]
# top100的分数
proposed_scores = scores[topk_indices]
  • 最终得到的药物中,目前已经处于临床实验的有6例,具体结果如下(排名,药物名称,相关分数)
[0]    Ribavirin          -0.21416784822940826
[4]    Dexamethasone      -0.9984006881713867
[8]    Colchicine         -1.080674648284912
[16]   Methylprednisolone -1.1618402004241943
[49]   Oseltamivir        -1.3885014057159424
[87]   Deferoxamine       -1.513066053390503
 

传送门:

github.com/gnn4dr/DRKG/

DRKG知识图谱:dgl-data.s3-us-west-2.amazonaws.com

DGL-KE:github.com/awslabs/dgl-

—完—

@量子位 · 追踪AI技术和产品新动态

深有感触的朋友,欢迎赞同、关注、分享三连վ'ᴗ' ի ❤

编辑于 06-10

文章被以下专栏收录