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 ingestiondeepchem — alternative ML framework for drug discovery (TensorFlow/PyTorch)scientific-skills:esm — ESM protein language models (direct HuggingFace usage)
1---2name: torchdrug3description: Use when working with TorchDrug for graph-based drug discovery and molecular ML. Covers molecular property prediction, protein modeling, knowledge graph reasoning, molecular generation, retrosynthesis, and GNN architectures on chemical data.4---56# TorchDrug78PyTorch toolkit for drug discovery — graph neural networks on molecules, proteins, and biomedical knowledge graphs. 40+ datasets, 20+ model architectures, modular task/model interface.910## When to Use This Skill1112- Predicting molecular properties (solubility, toxicity, BBB penetration, quantum chemistry)13- Protein function/stability/localization/interaction prediction14- Drug-target binding affinity (PDBBind, BindingDB)15- Knowledge graph completion and drug repurposing (Hetionet)16- De novo molecular generation and property optimization (GCPN, flows)17- Retrosynthesis planning (USPTO-50k, CenterIdentification + SynthonCompletion)18- Training GNNs (GCN, GAT, GIN, SchNet, GearNet, RGCN) on chemical data19- Transfer learning with pre-trained protein models (ESM, ProteinBERT)2021## Quick Start2223```python24from torchdrug import datasets, models, tasks25import torch26from torch.utils.data import DataLoader2728# 1. Dataset29dataset = datasets.BBBP("~/datasets/")30train_set, valid_set, test_set = dataset.split()3132# 2. Model33model = models.GIN(34 input_dim=dataset.node_feature_dim,35 hidden_dims=[256, 256, 256],36 edge_input_dim=dataset.edge_feature_dim,37 batch_norm=True, readout="mean"38)3940# 3. Task41task = tasks.PropertyPrediction(42 model, task=dataset.tasks,43 criterion="bce", metric=["auroc", "auprc"]44)4546# 4. Train47optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)48for epoch in range(100):49 for batch in DataLoader(train_set, batch_size=32, shuffle=True):50 loss = task(batch)51 optimizer.zero_grad(); loss.backward(); optimizer.step()52```5354## Router — What to Read5556| Task | Reference |57|------|-----------|58| Data structures (Graph, Molecule, Protein), training loop, task/model interface | `references/core-data.md` |59| Molecular property prediction: datasets, tasks, model selection, training | `references/molecular-property.md` |60| Protein modeling: sequence & structure models, datasets, pre-training | `references/protein-modeling.md` |61| Knowledge graph completion, drug repurposing, Hetionet | `references/knowledge-graphs.md` |62| Molecular generation: GCPN, flows, property optimization | `references/molecular-generation.md` |63| Retrosynthesis: CenterIdentification, SynthonCompletion, USPTO-50k | `references/retrosynthesis.md` |64| Full model catalog: GCN, GAT, GIN, SchNet, GearNet, ESM, TransE, RotatE… | `references/models-reference.md` |6566## Key Submodules6768| Module | Role |69|--------|------|70| `torchdrug.data` | Graph, Molecule, Protein, PackedGraph |71| `torchdrug.datasets` | 40+ curated datasets |72| `torchdrug.models` | GNN, protein, KG embedding, generative models |73| `torchdrug.tasks` | PropertyPrediction, KGCompletion, Generation, Retrosynthesis |74| `torchdrug.transforms` | VirtualNode, VirtualEdge, TruncateProtein |75| `torchdrug.layers` | MessagePassingBase and building blocks |76| `torchdrug.core` | Configurable, Registry (serialization) |7778## Installation7980```bash81pip install torchdrug # CPU / CUDA (uses system torch)82pip install torchdrug[full] # with optional extras83```8485```python86import torchdrug; print(torchdrug.__version__) # verify87```8889## Related Skills9091- `rdkit` — molecular I/O, fingerprints, conformers before TorchDrug ingestion92- `deepchem` — alternative ML framework for drug discovery (TensorFlow/PyTorch)93- `scientific-skills:esm` — ESM protein language models (direct HuggingFace usage)
Kdevos12