探索「老藥新用」最短路徑:亞馬遜AI Lab開源大規模藥物重定位知識圖譜DRKG

最近，來自亞馬遜上海 AI 實驗室、亞馬遜 AI 北美、明尼蘇達大學、俄亥俄州立大學、湖南大學等機構的團隊，正式開源了大規模藥物重定位知識圖譜 DRKG和一套完整的用於進行藥物重定位研究的機器學習工具，助力新冠及其他疾病的藥物重定位研究。

自 2019 年 12 月至今，新型冠狀病毒在全球迅速擴散已導致近 760 萬人感染，40 餘萬人死亡。目前急需快速有效的新冠病毒有效藥物的發現路徑。藥物重定位是一種將現有藥物用於治療新的適應症的藥物發現方式。相比較傳統的新藥開發，它可以有效縮短藥物研發周期，降低成本，規避風險，是一種非常有前景的新冠肺炎治療策略。近日，亞馬遜上海 AI 實驗室聯合來自亞馬遜 AI 北美、明尼蘇達大學、俄亥俄州立大學、湖南大學的研究者，共同構建了大規模藥物重定位知識圖譜 DRKG(Drug Repurposing Knowledge Graph)，以及一套完整的用於進行藥物重定位研究的機器學習工具，並將其在 github 上開源給全世界研究機構，幫助相關研究人員更有效地對新冠病毒及其它疾病（如阿爾茨海默病）進行藥物重定位研究。項目地址：https://github.com/gnn4dr/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 分鐘內完成訓練。項目地址：https://github.com/awslabs/dgl-ke

DRKG 在其開原始碼倉庫中提供了使用 DGL-KE 來生成 DRKG 知識圖譜中實體和關係的低緯嵌入向量表示的案例以及基於預訓練的 DRKG 知識圖譜嵌入向量表示預測新冠病毒候選藥物的案例。初步的實驗結果表明：使用此類機器學習工具進行冠狀病毒的藥物發現，能以較高的分數識別出目前正在進行臨床試驗的多種冠狀病毒候選藥物。
使用 DGL-KE 生成 DRKG 知識圖譜的低緯嵌入向量表示DGL-KE 是一個基於 DGL 圖深度學習框架(https://github.com/dmlc/dgl)，面向知識圖譜嵌入學習方法領域的高性能、高可擴展的開源軟體庫，DRKG 利用 DGL-KE 軟體包進行 DRKG 知識圖譜的低緯嵌入向量表示的生成。

import syssys.path.insert(1, '../utils')from utils import download_and_extractdownload_and_extract()drkg_file = '../data/drkg/drkg.tsv'

import pandas as pdimport 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]))
!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
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')
完整案例可以在此獲取：https://github.com/gnn4dr/DRKG/blob/master/embedding_analysis/Train_embeddings.ipynb 使用預訓練的 DRKG 知識圖譜的嵌入向量預測新冠病毒候選藥物在論文《Repurpose Open Data to Discover Therapeutics for COVID-19 using Deep Learning》中，亞馬遜 AI、湖南大學、克裡夫蘭診所勒納中心基因組醫學研究所、明尼蘇達大學的研究者提出了結合使用知識圖譜的嵌入和基因集富集分析的方法來進行新冠病毒老藥新用藥物的預測。DRKG 基於該論文的思路，提供了基於預訓練的 DRKG 知識圖譜嵌入向量表示進行新冠病毒候選藥物預測案例。論文地址：https://arxiv.org/abs/2005.10831首先我們將基於 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'])

# 讀取預訓練 embeddingentity_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]
其中，h  為 head 即藥物，r 為關係，t 為 tail 即病毒。gamma 為訓練中使用的常數。
import torch.nn.functional as fn
gamma=12.0def 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=100_, unique_indices = np.unique(dids, return_index=True)topk_indices = np.sort(unique_indices)[:topk]proposed_dids = dids[topk_indices]proposed_scores = scores[topk_indices]
[0]    Ribavirin          -0.21416784822940826[4]    Dexamethasone      -0.9984006881713867[8]    Colchicine         -1.080674648284912[16]   Methylprednisolone -1.1618402004241943[49]   Oseltamivir        -1.3885014057159424[87]   Deferoxamine       -1.513066053390503
完整案例可以在此處獲取：https://github.com/gnn4dr/DRKG/blob/master/drug_repurpose/COVID-19_drug_repurposing.ipynb

✄---