Skip to main content

A powerful and flexible machine learning platform for drug discovery

Project description


Open In Colab Contributions License Apache-2.0

Docs | Tutorials | Benchmarks | Papers Implemented

TorchDrug is a PyTorch-based machine learning toolbox designed for several purposes.

  • Easy implementation of graph operations in a PyTorchic style with GPU support
  • Being friendly to practitioners with minimal knowledge about drug discovery
  • Rapid prototyping of machine learning research


TorchDrug is compatible with Python 3.7/3.8 and PyTorch >= 1.4.0.

From Conda

conda install -c milagraph -c conda-forge torchdrug

From Pip

pip3 install torch==1.9.0
pip3 install torch-scatter -f
pip3 install torchdrug

To install torch-scatter for other PyTorch or CUDA versions, please see the instructions in

From Source

git clone
cd torchdrug
pip install -r requirements.txt
python install

Quick Start

TorchDrug is designed for humans and focused on graph structured data. It enables easy implementation of graph operations in machine learning models. All the operations in TorchDrug are backed by PyTorch framework, and support GPU acceleration and auto differentiation.

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

Molecules are also supported in TorchDrug. You can get the desired molecule properties without any domain knowledge.

mol = data.Molecule.from_smiles("CCOC(=O)N", node_feature="default", edge_feature="default")

You may also register custom node, edge or graph attributes. They will be automatically processed during indexing operations.

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)

TorchDrug provides a wide range of common datasets and building blocks for drug discovery. With minimal code, you can apply standard models to solve your own problem.

import torch
from torchdrug import datasets

dataset = datasets.Tox21()
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set =, 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)

Training and inference are accelerated by multiple CPUs or GPUs. This can be seamlessly switched in TorchDrug by just a line of code.

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


Everyone is welcome to contribute to the development of TorchDrug. Please refer to contributing guidelines for more details.


TorchDrug is released under Apache-2.0 License.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for torchdrug, version 0.1.2.post1
Filename, size File type Python version Upload date Hashes
Filename, size torchdrug-0.1.2.post1-py3-none-any.whl (191.5 kB) File type Wheel Python version py3 Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page