TorchDrug

PyTorch toolkit for drug discovery — graph neural networks on molecules, proteins, and biomedical knowledge graphs. 40+ datasets, 20+ model architectures, modular task/model interface.

When to Use This Skill

• Predicting molecular properties (solubility, toxicity, BBB penetration, quantum chemistry)
• Protein function/stability/localization/interaction prediction
• Drug-target binding affinity (PDBBind, BindingDB)
• Knowledge graph completion and drug repurposing (Hetionet)
• De novo molecular generation and property optimization (GCPN, flows)
• Retrosynthesis planning (USPTO-50k, CenterIdentification + SynthonCompletion)
• Training GNNs (GCN, GAT, GIN, SchNet, GearNet, RGCN) on chemical data
• Transfer learning with pre-trained protein models (ESM, ProteinBERT)

Quick Start

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

# 1. Dataset
dataset = datasets.BBBP("~/datasets/")
train_set, valid_set, test_set = dataset.split()

# 2. 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"
)

# 3. Task
task = tasks.PropertyPrediction(
    model, task=dataset.tasks,
    criterion="bce", metric=["auroc", "auprc"]
)

# 4. Train
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
for epoch in range(100):
    for batch in DataLoader(train_set, batch_size=32, shuffle=True):
        loss = task(batch)
        optimizer.zero_grad(); loss.backward(); optimizer.step()

Router — What to Read

Task	Reference
Data structures (Graph, Molecule, Protein), training loop, task/model interface	references/core-data.md
Molecular property prediction: datasets, tasks, model selection, training	references/molecular-property.md
Protein modeling: sequence & structure models, datasets, pre-training	references/protein-modeling.md
Knowledge graph completion, drug repurposing, Hetionet	references/knowledge-graphs.md
Molecular generation: GCPN, flows, property optimization	references/molecular-generation.md
Retrosynthesis: CenterIdentification, SynthonCompletion, USPTO-50k	references/retrosynthesis.md
Full model catalog: GCN, GAT, GIN, SchNet, GearNet, ESM, TransE, RotatE…	references/models-reference.md

Key Submodules

Module	Role
torchdrug.data	Graph, Molecule, Protein, PackedGraph
torchdrug.datasets	40+ curated datasets
torchdrug.models	GNN, protein, KG embedding, generative models
torchdrug.tasks	PropertyPrediction, KGCompletion, Generation, Retrosynthesis
torchdrug.transforms	VirtualNode, VirtualEdge, TruncateProtein
torchdrug.layers	MessagePassingBase and building blocks
torchdrug.core	Configurable, Registry (serialization)

Installation

pip install torchdrug          # CPU / CUDA (uses system torch)
pip install torchdrug[full]    # with optional extras

import torchdrug; print(torchdrug.__version__)  # verify

Related Skills

• rdkit — molecular I/O, fingerprints, conformers before TorchDrug ingestion
• deepchem — alternative ML framework for drug discovery (TensorFlow/PyTorch)
• scientific-skills:esm — ESM protein language models (direct HuggingFace usage)