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许可下发布。