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

本文經授權轉載自機器之心（almosthuman2014），未經授權禁止二次轉載與摘編。

本文長度爲3200字，建議閱讀9分鐘

本文帶你瞭解亞馬遜AI實驗室的開源機器學習工具。

[ 摘要 ]最近，來自亞馬遜上海 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 知識圖譜嵌入向量表示預測新冠病毒候選藥物的案例。

案例鏈接：

https://github.com/gnn4dr/DRKG/blob/master/embedding_analysis/Train_embeddings.ipynb

https://github.com/gnn4dr/DRKG/blob/master/drug_repurpose/COVID-19_drug_repurposing.ipynb

初步的實驗結果表明：使用此類機器學習工具進行冠狀病毒的藥物發現，能以較高的分數識別出目前正在進行臨牀試驗的多種冠狀病毒候選藥物。

使用 DGL-KE 生成 DRKG 知識圖譜的低緯嵌入向量表示

DGL-KE 是一個基於 DGL 圖深度學習框架(https://github.com/dmlc/dgl)，面向知識圖譜嵌入學習方法領域的高性能、高可擴展的開源軟件庫，DRKG 利用 DGL-KE 軟件包進行 DRKG 知識圖譜的低緯嵌入向量表示的生成。

• 下載 DRKG 知識圖譜，DRKG 知識圖譜已開源：

  https://dgl-data.s3-us-west-2.amazonaws.com/dataset/DRKG/drkg.tar.gz

import sys
sys.path.insert(1, '../utils')
from utils import 
download_and_extractdownload_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 並行進行訓練（使用其他知識圖譜嵌入算法以及其他機型可以參考 https://aws-dglke.readthedocs.io/en/latest/index.html 中的說明）。

!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')

完整案例可以在此獲取：

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'])

• 獲取預訓練 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

完整案例可以在此處獲取：

https://github.com/gnn4dr/DRKG/blob/master/drug_repurpose/COVID-19_drug_repurposing.ipynb

編輯：黃繼彥