TorchDrug

Overview

TorchDrug is a comprehensive PyTorch-based machine learning toolbox for drug discovery and molecular science. Apply graph neural networks, pre-trained models, and task definitions to molecules, proteins, and biological knowledge graphs, including molecular property prediction, protein modeling, knowledge graph reasoning, molecular generation, retrosynthesis planning, with 40+ curated datasets and 20+ model architectures.

When to Use This Skill

This skill should be used when working with:

Data Types:

• SMILES strings or molecular structures
• Protein sequences or 3D structures (PDB files)
• Chemical reactions and retrosynthesis
• Biomedical knowledge graphs
• Drug discovery datasets

Tasks:

• Predicting molecular properties (solubility, toxicity, activity)
• Protein function or structure prediction
• Drug-target binding prediction
• Generating new molecular structures
• Planning chemical synthesis routes
• Link prediction in biomedical knowledge bases
• Training graph neural networks on scientific data

Libraries and Integration:

• TorchDrug is the primary library
• Often used with RDKit for cheminformatics
• Compatible with PyTorch and PyTorch Lightning
• Integrates with AlphaFold and ESM for proteins

Getting Started

Installation

TorchDrug 0.2.1 (latest on PyPI, July 2023) requires Python 3.7–3.10 and PyTorch 1.8–2.0. Install PyTorch and torch-scatter / torch-cluster first (wheel URL depends on your PyTorch and CUDA versions — see installation docs).

uv pip install torch
# Match torch/CUDA in the URL, e.g. torch-2.0.0+cu118 or cpu
uv pip install torch-scatter torch-cluster -f https://pytorch-geometric.com/whl/torch-2.0.0+cu118.html
uv pip install torchdrug==0.2.1

On Apple Silicon, compile scatter/cluster from source; TorchDrug runs on CPU only (no MPS). Conda: conda install torchdrug -c milagraph -c conda-forge -c pytorch -c pyg.

Quick Example

import torch
from torchdrug import datasets, models, tasks
from torch.utils.data import DataLoader

# Load molecular dataset
dataset = datasets.BBBP("~/molecule-datasets/")
train_set, valid_set, test_set = dataset.split()

# Define GNN model
model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean"
)

# Create property prediction task
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=["auroc", "auprc"]
)

# Train with PyTorch
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        loss = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Core Capabilities

1. Molecular Property Prediction

Predict chemical, physical, and biological properties of molecules from structure.

Use Cases:

• Drug-likeness and ADMET properties
• Toxicity screening
• Quantum chemistry properties
• Binding affinity prediction

Key Components:

• 20+ molecular datasets (BBBP, HIV, Tox21, QM9, etc.)
• GNN models (GIN, GAT, SchNet)
• PropertyPrediction and MultipleBinaryClassification tasks

Reference: See references/molecular_property_prediction.md for:

• Complete dataset catalog
• Model selection guide
• Training workflows and best practices
• Feature engineering details

2. Protein Modeling

Work with protein sequences, structures, and properties.

Use Cases:

• Enzyme function prediction
• Protein stability and solubility
• Subcellular localization
• Protein-protein interactions
• Structure prediction

Key Components:

• 15+ protein datasets (EnzymeCommission, GeneOntology, PDBBind, etc.)
• Sequence models (ESM, ProteinBERT, ProteinLSTM)
• Structure models (GearNet, SchNet)
• Multiple task types for different prediction levels

Reference: See references/protein_modeling.md for:

• Protein-specific datasets
• Sequence vs structure models
• Pre-training strategies
• Integration with AlphaFold and ESM

3. Knowledge Graph Reasoning

Predict missing links and relationships in biological knowledge graphs.

Use Cases:

• Drug repurposing
• Disease mechanism discovery
• Gene-disease associations
• Multi-hop biomedical reasoning

Key Components:

• General KGs (FB15k, WN18) and biomedical (Hetionet)
• Embedding models (TransE, RotatE, ComplEx)
• KnowledgeGraphCompletion task

Reference: See references/knowledge_graphs.md for:

• Knowledge graph datasets (including Hetionet with 45k biomedical entities)
• Embedding model comparison
• Evaluation metrics and protocols
• Biomedical applications

4. Molecular Generation

Generate novel molecular structures with desired properties.

Use Cases:

• De novo drug design
• Lead optimization
• Chemical space exploration
• Property-guided generation

Key Components:

• Autoregressive generation
• GCPN (policy-based generation)
• GraphAutoregressiveFlow
• Property optimization workflows

Reference: See references/molecular_generation.md for:

• Generation strategies (unconditional, conditional, scaffold-based)
• Multi-objective optimization
• Validation and filtering
• Integration with property prediction

5. Retrosynthesis

Predict synthetic routes from target molecules to starting materials.

Use Cases:

• Synthesis planning
• Route optimization
• Synthetic accessibility assessment
• Multi-step planning

Key Components:

• USPTO-50k reaction dataset
• CenterIdentification (reaction center prediction)
• SynthonCompletion (reactant prediction)
• End-to-end Retrosynthesis pipeline

Reference: See references/retrosynthesis.md for:

• Task decomposition (center ID → synthon completion)
• Multi-step synthesis planning
• Commercial availability checking
• Integration with other retrosynthesis tools

6. Graph Neural Network Models

Comprehensive catalog of GNN architectures for different data types and tasks.

Available Models:

• General GNNs: GCN, GAT, GIN, RGCN, MPNN
• 3D-aware: SchNet, GearNet
• Protein-specific: ESM, ProteinBERT, GearNet
• Knowledge graph: TransE, RotatE, ComplEx, SimplE
• Generative: GraphAutoregressiveFlow

Reference: See references/models_architectures.md for:

• Detailed model descriptions
• Model selection guide by task and dataset
• Architecture comparisons
• Implementation tips

7. Datasets

40+ curated datasets spanning chemistry, biology, and knowledge graphs.

Categories:

• Molecular properties (drug discovery, quantum chemistry)
• Protein properties (function, structure, interactions)
• Knowledge graphs (general and biomedical)
• Retrosynthesis reactions

Reference: See references/datasets.md for:

• Complete dataset catalog with sizes and tasks
• Dataset selection guide
• Loading and preprocessing
• Splitting strategies (random, scaffold)

Common Workflows

Workflow 1: Molecular Property Prediction

Scenario: Predict blood-brain barrier penetration for drug candidates.

Steps:

1. Load dataset: datasets.BBBP()
2. Choose model: GIN for molecular graphs
3. Define task: PropertyPrediction with binary classification
4. Train with scaffold split for realistic evaluation
5. Evaluate using AUROC and AUPRC

Navigation: references/molecular_property_prediction.md → Dataset selection → Model selection → Training

Workflow 2: Protein Function Prediction

Scenario: Predict enzyme function from sequence.

Steps:

1. Load dataset: datasets.EnzymeCommission()
2. Choose model: ESM (pre-trained) or GearNet (with structure)
3. Define task: PropertyPrediction with multi-class classification
4. Fine-tune pre-trained model or train from scratch
5. Evaluate using accuracy and per-class metrics

Navigation: references/protein_modeling.md → Model selection (sequence vs