torchdrug

和

文檔|教程|基準|已實施論文

Torchdrug是一種基於Pytorch的機器學習工具箱，專為多種目的而設計。

• 輕鬆地以pytorchic風格實施具有GPU支持的圖形操作
• 對對藥物發現知識了解最少的從業者友好
• 機器學習研究的快速原型製作

安裝

Torchdrug可以安裝在Linux，Windows或MacOS上。它與3.7 <= python <= 3.10和pytorch> = 1.8.0兼容。

來自康達

conda install torchdrug -c milagraph -c conda-forge -c pytorch -c pyg

來自pip

pip install torch==1.9.0
pip install torch-scatter torch-cluster -f https://pytorch-geometric.com/whl/torch-1.9.0+cu102.html
pip install torchdrug

要安裝其他Pytorch或CUDA版本的torch-scatter ，請在https://github.com/rusty1s/pytorch_scatter中查看說明

來自來源

git clone https://github.com/DeepGraphLearning/torchdrug
cd torchdrug
pip install -r requirements.txt
python setup.py install

Windows（PowerShell）

我們需要首先安裝Visual Studio的構建工具。然後，我們將以下模塊安裝在PowerShell中。

 Install-Module Pscx - AllowClobber
Install-Module VSSetup

使用以下命令在PowerShell中初始化Visual Studio。我們可以通過將其寫入PowerShell配置文件來為所有PowerShell會話設置此功能。根據您自己的情況更改圖書館路徑。

 Import-VisualStudioVars - Architecture x64
$ env: LIB += " ;C:Program FilesPython37libs "

蘋果矽（M1/M2芯片）

我們需要pytorch> = 1.13才能在蘋果矽上運行Torchdrug。對於torch-scatter和torch-cluster ，可以從他們的來源中彙編。注意Torchdrug不支持mps設備。

pip install torch==1.13.0
pip install git+https://github.com/rusty1s/pytorch_scatter.git
pip install git+https://github.com/rusty1s/pytorch_cluster.git
pip install torchdrug

快速開始

Torchdrug是為人類設計的，專注於圖形結構化數據。它可以輕鬆實現機器學習模型中的圖形操作。 Torchdrug中的所有操作均由Pytorch框架支持，並支持GPU加速和自動差異化。

 from torchdrug import data

edge_list = [[ 0 , 1 ], [ 1 , 2 ], [ 2 , 3 ], [ 3 , 4 ], [ 4 , 5 ], [ 5 , 0 ]]
graph = data . Graph ( edge_list , num_node = 6 )
graph = graph . cuda ()
# the subgraph induced by nodes 2, 3 & 4
subgraph = graph . subgraph ([ 2 , 3 , 4 ])

分子也得到了毛刺的支持。您可以在沒有任何領域知識的情況下獲得所需的分子屬性。

 mol = data . Molecule . from_smiles ( "CCOC(=O)N" , atom_feature = "default" , bond_feature = "default" )
print ( mol . node_feature )
print ( mol . atom_type )
print ( mol . to_scaffold ())

您也可以註冊自定義節點，邊緣或圖形屬性。它們將在索引操作期間自動處理。

 with mol . edge ():
	mol . is_CC_bond = ( mol . edge_list [:, : 2 ] == td . CARBON ). all ( dim = - 1 )
sub_mol = mol . subgraph ( mol . atom_type != td . NITROGEN )
print ( sub_mol . is_CC_bond )

Torchdrug提供了廣泛的常見數據集和用於藥物發現的基礎。使用最小代碼，您可以應用標準模型來解決自己的問題。

 import torch
from torchdrug import datasets

dataset = datasets . Tox21 ()
dataset [ 0 ]. visualize ()
lengths = [ int ( 0.8 * len ( dataset )), int ( 0.1 * len ( dataset ))]
lengths += [ len ( dataset ) - sum ( lengths )]
train_set , valid_set , test_set = torch . utils . data . random_split ( dataset , lengths )

 from torchdrug import models , tasks

model = models . GIN ( dataset . node_feature_dim , hidden_dims = [ 256 , 256 , 256 , 256 ])
task = tasks . PropertyPrediction ( model , task = dataset . tasks )

培訓和推理通過多個CPU或GPU加速。只需一系列代碼即可在Torchdrug中無縫切換。

 from torchdrug import core

# Single CPU / Multiple CPUs / Distributed CPUs
solver = core . Engine ( task , train_set , valid_set , test_set , optimizer )
# Single GPU
solver = core . Engine ( task , train_set , valid_set , test_set , optimizer , gpus = [ 0 ])
# Multiple GPUs
solver = core . Engine ( task , train_set , valid_set , test_set , optimizer , gpus = [ 0 , 1 , 2 , 3 ])
# Distributed GPUs
solver = core . Engine ( task , train_set , valid_set , test_set , optimizer , gpus = [ 0 , 1 , 2 , 3 , 0 , 1 , 2 , 3 ])

可以通過重量和偏見平台輕鬆地跟踪和管理實驗。

 solver = core . Engine ( task , train_set , valid_set , test_set , optimizer , logger = "wandb" )

貢獻

歡迎每個人為Torchdrug的發展做出貢獻。有關更多詳細信息，請參考貢獻指南。

執照

Torchdrug在Apache-2.0許可下發布。