torchdrug

torchdrug

Overview

TorchDrug is a comprehensive machine learning framework for drug discovery built on PyTorch. It provides graph-based molecular representations (atoms as nodes, bonds as edges), a library of graph neural network (GNN) architectures, benchmark datasets, and pretrained models for tasks including molecular property prediction, drug-target interaction, retrosynthesis, and generative molecular design. TorchDrug integrates with PyTorch Lightning and standard ML tooling, making it accessible to both computational chemists and ML practitioners.

When to Use

• Molecular property prediction: Training or fine-tuning GNN models to predict ADMET properties (solubility, toxicity, permeability) or bioactivity (IC50, Ki) from molecular graphs.
• Drug-target interaction (DTI) prediction: Building models that predict binding affinity between a compound (SMILES) and a protein (sequence or structure).
• Retrosynthesis prediction: Identifying plausible synthetic routes for a target molecule using template-based or template-free models.
• Pretraining on large molecular datasets: Leveraging pretrained GNN representations on ChEMBL or ZINC for transfer learning to small datasets.
• Molecular generation: Training graph-based generative models (GCPN, GraphAF) to design novel molecules with desired properties.
• Benchmarking GNN architectures: Comparing GraphConv, MPNN, GAT, AttentiveFP on standard MoleculeNet tasks.
• For fast fingerprint-based property prediction without deep learning, use RDKit + scikit-learn instead.
• For protein structure tasks (folding, docking), use ESMFold or DiffDock rather than TorchDrug.

Prerequisites

• Python packages: torchdrug, torch, torch-geometric, rdkit
• Environment: Python 3.8+, CUDA-compatible GPU recommended for training
• Data requirements: SMILES strings or molecular SDF files; protein sequences for DTI tasks

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118
pip install torch-geometric
pip install torchdrug
pip install rdkit

Quick Start

import torch
from torchdrug import data, datasets, models, tasks, core

# Load a benchmark dataset and train a GNN for property prediction
dataset = datasets.BBBP("~/data/bbbp", node_feature="default", edge_feature="default")
print(f"Dataset: {len(dataset)} molecules, task: BBBP (blood-brain barrier penetration)")