torchdrug

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:

Tasks:

Libraries and Integration:

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:

Key Components:

Reference: See references/molecular_property_prediction.md for:

2. Protein Modeling

Work with protein sequences, structures, and properties.

Use Cases:

Key Components:

Reference: See references/protein_modeling.md for:

3. Knowledge Graph Reasoning

Predict missing links and relationships in biological knowledge graphs.

Use Cases:

Key Components:

Reference: See references/knowledge_graphs.md for:

4. Molecular Generation

Generate novel molecular structures with desired properties.

Use Cases:

Key Components:

Reference: See references/molecular_generation.md for:

5. Retrosynthesis

Predict synthetic routes from target molecules to starting materials.

Use Cases:

Key Components:

Reference: See references/retrosynthesis.md for:

6. Graph Neural Network Models

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

Available Models:

Reference: See references/models_architectures.md for:

7. Datasets

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

Categories:

Reference: See references/datasets.md for:

Common Workflows

Workflow 1: Molecular Property Prediction

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

Steps:

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

Workflow 2: Protein Function Prediction

Scenario: Predict enzyme function from sequence.

Steps:

Navigation: references/protein_modeling.md → Model selection (sequence vs structure) → Pre-training strategies

Workflow 3: Drug Repurposing via Knowledge Graphs

Scenario: Find new disease treatments in Hetionet.

Steps:

Navigation: references/knowledge_graphs.md → Hetionet dataset → Model selection → Biomedical applications

Workflow 4: De Novo Molecule Generation

Scenario: Generate drug-like molecules optimized for target binding.

Steps:

Navigation: references/molecular_generation.md → Conditional generation → Multi-objective optimization

Workflow 5: Retrosynthesis Planning

Scenario: Plan synthesis route for target molecule.

Steps:

Navigation: references/retrosynthesis.md → Task types → Multi-step planning

Integration Patterns

With RDKit

Convert between TorchDrug molecules and RDKit:

from torchdrug import data
from rdkit import Chem

# SMILES → TorchDrug molecule
smiles = "CCO"
mol = data.Molecule.from_smiles(smiles)

# TorchDrug → RDKit
rdkit_mol = mol.to_molecule()

# RDKit → TorchDrug
rdkit_mol = Chem.MolFromSmiles(smiles)
mol = data.Molecule.from_molecule(rdkit_mol)

With AlphaFold/ESM

Use predicted structures:

from torchdrug import data

# Load AlphaFold predicted structure
protein = data.Protein.from_pdb("AF-P12345-F1-model_v4.pdb")

# Build graph with spatial edges
graph = protein.residue_graph(
    node_position="ca",
    edge_types=["sequential", "radius"],
    radius_cutoff=10.0
)

With PyTorch Lightning

Wrap tasks for Lightning training:

import pytorch_lightning as pl

class LightningTask(pl.LightningModule):
    def __init__(self, torchdrug_task):
        super().__init__()
        self.task = torchdrug_task

    def training_step(self, batch, batch_idx):
        return self.task(batch)

    def validation_step(self, batch, batch_idx):
        pred = self.task.predict(batch)
        target = self.task.target(batch)
        return {"pred": pred, "target": target}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

Technical Details

For deep dives into TorchDrug's architecture:

Core Concepts: See references/core_concepts.md for:

Quick Reference Cheat Sheet

Choose Dataset:

Choose Model:

Common Tasks:

Understand Architecture:

Troubleshooting Common Issues

Issue: Dimension mismatch errors
→ Check model.input_dim matches dataset.node_feature_dim
→ See references/core_concepts.md → Essential Attributes

Issue: Poor performance on molecular tasks
→ Use scaffold splitting, not random
→ Try GIN instead of GCN
→ See references/molecular_property_prediction.md → Best Practices

Issue: Protein model not learning
→ Use pre-trained ESM for sequence tasks
→ Check edge construction for structure models
→ See references/protein_modeling.md → Training Workflows

Issue: Memory errors with large graphs
→ Reduce batch size
→ Use gradient accumulation
→ See references/core_concepts.md → Memory Efficiency

Issue: Generated molecules are invalid
→ Add validity constraints
→ Post-process with RDKit validation
→ See references/molecular_generation.md → Validation and Filtering

Version Notes (0.2.1)

Resources

Official Documentation: https://torchdrug.ai/docs/ (0.2.1)
GitHub: https://github.com/DeepGraphLearning/torchdrug
Paper: TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery

Summary

Navigate to the appropriate reference file based on your task:

Each reference provides comprehensive coverage of its domain with examples, best practices, and common use cases.