
Torchdrug
- 38 installs
- 16 repo stars
- Updated November 20, 2025
- jackspace/claudeskillz
Run PyTorch graph-neural-network workflows for drug discovery: molecular property prediction, protein modeling, molecular generation, and retrosynthesis.
About
TorchDrug is a PyTorch machine-learning toolbox for drug discovery that applies GNNs and pre-trained models to molecules, proteins, and biomedical knowledge graphs. A developer uses it to predict molecular properties, generate molecules, plan retrosynthesis, or reason over biomedical graphs with 40+ curated datasets.
- 20+ model architectures (GIN, GAT, SchNet, ESM) across property prediction, generation, and KG reasoning
- Integrates with RDKit, AlphaFold/ESM, and PyTorch Lightning
Torchdrug by the numbers
- 38 all-time installs (skills.sh)
- Ranked #1,015 of 2,065 Data Science & ML skills by installs in the Skillselion catalog
- Data as of Aug 2, 2026 (Skillselion catalog sync)
npx skills add https://github.com/jackspace/claudeskillz --skill torchdrugAdd your badge
Show developers this skill is listed on Skillselion. Paste this into your README.
| Installs | 38 |
|---|---|
| repo stars | ★ 16 |
| Last updated | November 20, 2025 |
| Repository | jackspace/claudeskillz ↗ |
What it does
Run PyTorch graph-neural-network workflows for drug discovery: molecular property prediction, protein modeling, molecular generation, and retrosynthesis.
Files
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
pip install torchdrug
# Or with optional dependencies
pip install torchdrug[full]Quick Example
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 structure) → Pre-training strategies
Workflow 3: Drug Repurposing via Knowledge Graphs
Scenario: Find new disease treatments in Hetionet.
Steps: 1. Load dataset: datasets.Hetionet() 2. Choose model: RotatE or ComplEx 3. Define task: KnowledgeGraphCompletion 4. Train with negative sampling 5. Query for "Compound-treats-Disease" predictions 6. Filter by plausibility and mechanism
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: 1. Train property predictor on activity data 2. Choose generation approach: GCPN for RL-based optimization 3. Define reward function combining affinity, drug-likeness, synthesizability 4. Generate candidates with property constraints 5. Validate chemistry and filter by drug-likeness 6. Rank by multi-objective scoring
Navigation: references/molecular_generation.md → Conditional generation → Multi-objective optimization
Workflow 5: Retrosynthesis Planning
Scenario: Plan synthesis route for target molecule.
Steps: 1. Load dataset: datasets.USPTO50k() 2. Train center identification model (RGCN) 3. Train synthon completion model (GIN) 4. Combine into end-to-end retrosynthesis pipeline 5. Apply recursively for multi-step planning 6. Check commercial availability of building blocks
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:
- Architecture philosophy (modular, configurable)
- Data structures (Graph, Molecule, Protein, PackedGraph)
- Model interface and forward function signature
- Task interface (predict, target, forward, evaluate)
- Training workflows and best practices
- Loss functions and metrics
- Common pitfalls and debugging
Quick Reference Cheat Sheet
Choose Dataset:
- Molecular property →
references/datasets.md→ Molecular section - Protein task →
references/datasets.md→ Protein section - Knowledge graph →
references/datasets.md→ Knowledge graph section
Choose Model:
- Molecules →
references/models_architectures.md→ GNN section → GIN/GAT/SchNet - Proteins (sequence) →
references/models_architectures.md→ Protein section → ESM - Proteins (structure) →
references/models_architectures.md→ Protein section → GearNet - Knowledge graph →
references/models_architectures.md→ KG section → RotatE/ComplEx
Common Tasks:
- Property prediction →
references/molecular_property_prediction.mdorreferences/protein_modeling.md - Generation →
references/molecular_generation.md - Retrosynthesis →
references/retrosynthesis.md - KG reasoning →
references/knowledge_graphs.md
Understand Architecture:
- Data structures →
references/core_concepts.md→ Data Structures - Model design →
references/core_concepts.md→ Model Interface - Task design →
references/core_concepts.md→ Task Interface
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
Resources
Official Documentation: https://torchdrug.ai/docs/ 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:
1. Molecular property prediction → molecular_property_prediction.md 2. Protein modeling → protein_modeling.md 3. Knowledge graphs → knowledge_graphs.md 4. Molecular generation → molecular_generation.md 5. Retrosynthesis → retrosynthesis.md 6. Model selection → models_architectures.md 7. Dataset selection → datasets.md 8. Technical details → core_concepts.md
Each reference provides comprehensive coverage of its domain with examples, best practices, and common use cases.
{
"description": "\"Graph-based drug discovery toolkit. Molecular property prediction (ADMET), protein modeling, knowledge graph reasoning, molecular generation, retrosynthesis, GNNs (GIN, GAT, SchNet), 40+ datasets, for PyTorch-based ML on molecules, proteins, and biomedical graphs.\"",
"references": {
"files": [
"references/core_concepts.md",
"references/datasets.md",
"references/knowledge_graphs.md",
"references/models_architectures.md",
"references/molecular_generation.md",
"references/molecular_property_prediction.md",
"references/protein_modeling.md",
"references/retrosynthesis.md"
]
},
"content": "### Installation\r\n\r\n```bash\r\npip install torchdrug\r\npip install torchdrug[full]\r\n```\r\n\r\n### Quick Example\r\n\r\n```python\r\nfrom torchdrug import datasets, models, tasks\r\nfrom torch.utils.data import DataLoader\r\n\r\ndataset = datasets.BBBP(\"~/molecule-datasets/\")\r\ntrain_set, valid_set, test_set = dataset.split()\r\n\r\nmodel = models.GIN(\r\n input_dim=dataset.node_feature_dim,\r\n hidden_dims=[256, 256, 256],\r\n edge_input_dim=dataset.edge_feature_dim,\r\n batch_norm=True,\r\n readout=\"mean\"\r\n)\r\n\r\ntask = tasks.PropertyPrediction(\r\n model,\r\n task=dataset.tasks,\r\n criterion=\"bce\",\r\n metric=[\"auroc\", \"auprc\"]\r\n)\r\n\r\n\r\n### With RDKit\r\n\r\nConvert between TorchDrug molecules and RDKit:\r\n```python\r\nfrom torchdrug import data\r\nfrom rdkit import Chem\r\n\r\nsmiles = \"CCO\"\r\nmol = data.Molecule.from_smiles(smiles)\r\n\r\nrdkit_mol = mol.to_molecule()\r\n\r\nrdkit_mol = Chem.MolFromSmiles(smiles)\r\nmol = data.Molecule.from_molecule(rdkit_mol)\r\n```\r\n\r\n### With AlphaFold/ESM\r\n\r\nUse predicted structures:\r\n```python\r\nfrom torchdrug import data\r\n\r\nprotein = data.Protein.from_pdb(\"AF-P12345-F1-model_v4.pdb\")",
"name": "torchdrug",
"id": "scientific-pkg-torchdrug",
"sections": {
"Quick Reference Cheat Sheet": "**Choose Dataset:**\r\n- Molecular property → `references/datasets.md` → Molecular section\r\n- Protein task → `references/datasets.md` → Protein section\r\n- Knowledge graph → `references/datasets.md` → Knowledge graph section\r\n\r\n**Choose Model:**\r\n- Molecules → `references/models_architectures.md` → GNN section → GIN/GAT/SchNet\r\n- Proteins (sequence) → `references/models_architectures.md` → Protein section → ESM\r\n- Proteins (structure) → `references/models_architectures.md` → Protein section → GearNet\r\n- Knowledge graph → `references/models_architectures.md` → KG section → RotatE/ComplEx\r\n\r\n**Common Tasks:**\r\n- Property prediction → `references/molecular_property_prediction.md` or `references/protein_modeling.md`\r\n- Generation → `references/molecular_generation.md`\r\n- Retrosynthesis → `references/retrosynthesis.md`\r\n- KG reasoning → `references/knowledge_graphs.md`\r\n\r\n**Understand Architecture:**\r\n- Data structures → `references/core_concepts.md` → Data Structures\r\n- Model design → `references/core_concepts.md` → Model Interface\r\n- Task design → `references/core_concepts.md` → Task Interface",
"Technical Details": "For deep dives into TorchDrug's architecture:\r\n\r\n**Core Concepts:** See `references/core_concepts.md` for:\r\n- Architecture philosophy (modular, configurable)\r\n- Data structures (Graph, Molecule, Protein, PackedGraph)\r\n- Model interface and forward function signature\r\n- Task interface (predict, target, forward, evaluate)\r\n- Training workflows and best practices\r\n- Loss functions and metrics\r\n- Common pitfalls and debugging",
"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.",
"Troubleshooting Common Issues": "**Issue: Dimension mismatch errors**\r\n→ Check `model.input_dim` matches `dataset.node_feature_dim`\r\n→ See `references/core_concepts.md` → Essential Attributes\r\n\r\n**Issue: Poor performance on molecular tasks**\r\n→ Use scaffold splitting, not random\r\n→ Try GIN instead of GCN\r\n→ See `references/molecular_property_prediction.md` → Best Practices\r\n\r\n**Issue: Protein model not learning**\r\n→ Use pre-trained ESM for sequence tasks\r\n→ Check edge construction for structure models\r\n→ See `references/protein_modeling.md` → Training Workflows\r\n\r\n**Issue: Memory errors with large graphs**\r\n→ Reduce batch size\r\n→ Use gradient accumulation\r\n→ See `references/core_concepts.md` → Memory Efficiency\r\n\r\n**Issue: Generated molecules are invalid**\r\n→ Add validity constraints\r\n→ Post-process with RDKit validation\r\n→ See `references/molecular_generation.md` → Validation and Filtering",
"Integration Patterns": "graph = protein.residue_graph(\r\n node_position=\"ca\",\r\n edge_types=[\"sequential\", \"radius\"],\r\n radius_cutoff=10.0\r\n)\r\n```\r\n\r\n### With PyTorch Lightning\r\n\r\nWrap tasks for Lightning training:\r\n```python\r\nimport pytorch_lightning as pl\r\n\r\nclass LightningTask(pl.LightningModule):\r\n def __init__(self, torchdrug_task):\r\n super().__init__()\r\n self.task = torchdrug_task\r\n\r\n def training_step(self, batch, batch_idx):\r\n return self.task(batch)\r\n\r\n def validation_step(self, batch, batch_idx):\r\n pred = self.task.predict(batch)\r\n target = self.task.target(batch)\r\n return {\"pred\": pred, \"target\": target}\r\n\r\n def configure_optimizers(self):\r\n return torch.optim.Adam(self.parameters(), lr=1e-3)\r\n```",
"When to Use This Skill": "This skill should be used when working with:\r\n\r\n**Data Types:**\r\n- SMILES strings or molecular structures\r\n- Protein sequences or 3D structures (PDB files)\r\n- Chemical reactions and retrosynthesis\r\n- Biomedical knowledge graphs\r\n- Drug discovery datasets\r\n\r\n**Tasks:**\r\n- Predicting molecular properties (solubility, toxicity, activity)\r\n- Protein function or structure prediction\r\n- Drug-target binding prediction\r\n- Generating new molecular structures\r\n- Planning chemical synthesis routes\r\n- Link prediction in biomedical knowledge bases\r\n- Training graph neural networks on scientific data\r\n\r\n**Libraries and Integration:**\r\n- TorchDrug is the primary library\r\n- Often used with RDKit for cheminformatics\r\n- Compatible with PyTorch and PyTorch Lightning\r\n- Integrates with AlphaFold and ESM for proteins",
"Summary": "Navigate to the appropriate reference file based on your task:\r\n\r\n1. **Molecular property prediction** → `molecular_property_prediction.md`\r\n2. **Protein modeling** → `protein_modeling.md`\r\n3. **Knowledge graphs** → `knowledge_graphs.md`\r\n4. **Molecular generation** → `molecular_generation.md`\r\n5. **Retrosynthesis** → `retrosynthesis.md`\r\n6. **Model selection** → `models_architectures.md`\r\n7. **Dataset selection** → `datasets.md`\r\n8. **Technical details** → `core_concepts.md`\r\n\r\nEach reference provides comprehensive coverage of its domain with examples, best practices, and common use cases.",
"Getting Started": "optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)\r\ntrain_loader = DataLoader(train_set, batch_size=32, shuffle=True)\r\n\r\nfor epoch in range(100):\r\n for batch in train_loader:\r\n loss = task(batch)\r\n optimizer.zero_grad()\r\n loss.backward()\r\n optimizer.step()\r\n```",
"Core Capabilities": "### 1. Molecular Property Prediction\r\n\r\nPredict chemical, physical, and biological properties of molecules from structure.\r\n\r\n**Use Cases:**\r\n- Drug-likeness and ADMET properties\r\n- Toxicity screening\r\n- Quantum chemistry properties\r\n- Binding affinity prediction\r\n\r\n**Key Components:**\r\n- 20+ molecular datasets (BBBP, HIV, Tox21, QM9, etc.)\r\n- GNN models (GIN, GAT, SchNet)\r\n- PropertyPrediction and MultipleBinaryClassification tasks\r\n\r\n**Reference:** See `references/molecular_property_prediction.md` for:\r\n- Complete dataset catalog\r\n- Model selection guide\r\n- Training workflows and best practices\r\n- Feature engineering details\r\n\r\n### 2. Protein Modeling\r\n\r\nWork with protein sequences, structures, and properties.\r\n\r\n**Use Cases:**\r\n- Enzyme function prediction\r\n- Protein stability and solubility\r\n- Subcellular localization\r\n- Protein-protein interactions\r\n- Structure prediction\r\n\r\n**Key Components:**\r\n- 15+ protein datasets (EnzymeCommission, GeneOntology, PDBBind, etc.)\r\n- Sequence models (ESM, ProteinBERT, ProteinLSTM)\r\n- Structure models (GearNet, SchNet)\r\n- Multiple task types for different prediction levels\r\n\r\n**Reference:** See `references/protein_modeling.md` for:\r\n- Protein-specific datasets\r\n- Sequence vs structure models\r\n- Pre-training strategies\r\n- Integration with AlphaFold and ESM\r\n\r\n### 3. Knowledge Graph Reasoning\r\n\r\nPredict missing links and relationships in biological knowledge graphs.\r\n\r\n**Use Cases:**\r\n- Drug repurposing\r\n- Disease mechanism discovery\r\n- Gene-disease associations\r\n- Multi-hop biomedical reasoning\r\n\r\n**Key Components:**\r\n- General KGs (FB15k, WN18) and biomedical (Hetionet)\r\n- Embedding models (TransE, RotatE, ComplEx)\r\n- KnowledgeGraphCompletion task\r\n\r\n**Reference:** See `references/knowledge_graphs.md` for:\r\n- Knowledge graph datasets (including Hetionet with 45k biomedical entities)\r\n- Embedding model comparison\r\n- Evaluation metrics and protocols\r\n- Biomedical applications\r\n\r\n### 4. Molecular Generation\r\n\r\nGenerate novel molecular structures with desired properties.\r\n\r\n**Use Cases:**\r\n- De novo drug design\r\n- Lead optimization\r\n- Chemical space exploration\r\n- Property-guided generation\r\n\r\n**Key Components:**\r\n- Autoregressive generation\r\n- GCPN (policy-based generation)\r\n- GraphAutoregressiveFlow\r\n- Property optimization workflows\r\n\r\n**Reference:** See `references/molecular_generation.md` for:\r\n- Generation strategies (unconditional, conditional, scaffold-based)\r\n- Multi-objective optimization\r\n- Validation and filtering\r\n- Integration with property prediction\r\n\r\n### 5. Retrosynthesis\r\n\r\nPredict synthetic routes from target molecules to starting materials.\r\n\r\n**Use Cases:**\r\n- Synthesis planning\r\n- Route optimization\r\n- Synthetic accessibility assessment\r\n- Multi-step planning\r\n\r\n**Key Components:**\r\n- USPTO-50k reaction dataset\r\n- CenterIdentification (reaction center prediction)\r\n- SynthonCompletion (reactant prediction)\r\n- End-to-end Retrosynthesis pipeline\r\n\r\n**Reference:** See `references/retrosynthesis.md` for:\r\n- Task decomposition (center ID → synthon completion)\r\n- Multi-step synthesis planning\r\n- Commercial availability checking\r\n- Integration with other retrosynthesis tools\r\n\r\n### 6. Graph Neural Network Models\r\n\r\nComprehensive catalog of GNN architectures for different data types and tasks.\r\n\r\n**Available Models:**\r\n- General GNNs: GCN, GAT, GIN, RGCN, MPNN\r\n- 3D-aware: SchNet, GearNet\r\n- Protein-specific: ESM, ProteinBERT, GearNet\r\n- Knowledge graph: TransE, RotatE, ComplEx, SimplE\r\n- Generative: GraphAutoregressiveFlow\r\n\r\n**Reference:** See `references/models_architectures.md` for:\r\n- Detailed model descriptions\r\n- Model selection guide by task and dataset\r\n- Architecture comparisons\r\n- Implementation tips\r\n\r\n### 7. Datasets\r\n\r\n40+ curated datasets spanning chemistry, biology, and knowledge graphs.\r\n\r\n**Categories:**\r\n- Molecular properties (drug discovery, quantum chemistry)\r\n- Protein properties (function, structure, interactions)\r\n- Knowledge graphs (general and biomedical)\r\n- Retrosynthesis reactions\r\n\r\n**Reference:** See `references/datasets.md` for:\r\n- Complete dataset catalog with sizes and tasks\r\n- Dataset selection guide\r\n- Loading and preprocessing\r\n- Splitting strategies (random, scaffold)",
"Common Workflows": "### Workflow 1: Molecular Property Prediction\r\n\r\n**Scenario:** Predict blood-brain barrier penetration for drug candidates.\r\n\r\n**Steps:**\r\n1. Load dataset: `datasets.BBBP()`\r\n2. Choose model: GIN for molecular graphs\r\n3. Define task: `PropertyPrediction` with binary classification\r\n4. Train with scaffold split for realistic evaluation\r\n5. Evaluate using AUROC and AUPRC\r\n\r\n**Navigation:** `references/molecular_property_prediction.md` → Dataset selection → Model selection → Training\r\n\r\n### Workflow 2: Protein Function Prediction\r\n\r\n**Scenario:** Predict enzyme function from sequence.\r\n\r\n**Steps:**\r\n1. Load dataset: `datasets.EnzymeCommission()`\r\n2. Choose model: ESM (pre-trained) or GearNet (with structure)\r\n3. Define task: `PropertyPrediction` with multi-class classification\r\n4. Fine-tune pre-trained model or train from scratch\r\n5. Evaluate using accuracy and per-class metrics\r\n\r\n**Navigation:** `references/protein_modeling.md` → Model selection (sequence vs structure) → Pre-training strategies\r\n\r\n### Workflow 3: Drug Repurposing via Knowledge Graphs\r\n\r\n**Scenario:** Find new disease treatments in Hetionet.\r\n\r\n**Steps:**\r\n1. Load dataset: `datasets.Hetionet()`\r\n2. Choose model: RotatE or ComplEx\r\n3. Define task: `KnowledgeGraphCompletion`\r\n4. Train with negative sampling\r\n5. Query for \"Compound-treats-Disease\" predictions\r\n6. Filter by plausibility and mechanism\r\n\r\n**Navigation:** `references/knowledge_graphs.md` → Hetionet dataset → Model selection → Biomedical applications\r\n\r\n### Workflow 4: De Novo Molecule Generation\r\n\r\n**Scenario:** Generate drug-like molecules optimized for target binding.\r\n\r\n**Steps:**\r\n1. Train property predictor on activity data\r\n2. Choose generation approach: GCPN for RL-based optimization\r\n3. Define reward function combining affinity, drug-likeness, synthesizability\r\n4. Generate candidates with property constraints\r\n5. Validate chemistry and filter by drug-likeness\r\n6. Rank by multi-objective scoring\r\n\r\n**Navigation:** `references/molecular_generation.md` → Conditional generation → Multi-objective optimization\r\n\r\n### Workflow 5: Retrosynthesis Planning\r\n\r\n**Scenario:** Plan synthesis route for target molecule.\r\n\r\n**Steps:**\r\n1. Load dataset: `datasets.USPTO50k()`\r\n2. Train center identification model (RGCN)\r\n3. Train synthon completion model (GIN)\r\n4. Combine into end-to-end retrosynthesis pipeline\r\n5. Apply recursively for multi-step planning\r\n6. Check commercial availability of building blocks\r\n\r\n**Navigation:** `references/retrosynthesis.md` → Task types → Multi-step planning",
"Resources": "**Official Documentation:** https://torchdrug.ai/docs/\r\n**GitHub:** https://github.com/DeepGraphLearning/torchdrug\r\n**Paper:** TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery"
}
}---
name: torchdrug
description: "Graph-based drug discovery toolkit. Molecular property prediction (ADMET), protein modeling, knowledge graph reasoning, molecular generation, retrosynthesis, GNNs (GIN, GAT, SchNet), 40+ datasets, for PyTorch-based ML on molecules, proteins, and biomedical graphs."
---
# 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
```bash
pip install torchdrug
# Or with optional dependencies
pip install torchdrug[full]
```
### Quick Example
```python
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 structure) → Pre-training strategies
### Workflow 3: Drug Repurposing via Knowledge Graphs
**Scenario:** Find new disease treatments in Hetionet.
**Steps:**
1. Load dataset: `datasets.Hetionet()`
2. Choose model: RotatE or ComplEx
3. Define task: `KnowledgeGraphCompletion`
4. Train with negative sampling
5. Query for "Compound-treats-Disease" predictions
6. Filter by plausibility and mechanism
**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:**
1. Train property predictor on activity data
2. Choose generation approach: GCPN for RL-based optimization
3. Define reward function combining affinity, drug-likeness, synthesizability
4. Generate candidates with property constraints
5. Validate chemistry and filter by drug-likeness
6. Rank by multi-objective scoring
**Navigation:** `references/molecular_generation.md` → Conditional generation → Multi-objective optimization
### Workflow 5: Retrosynthesis Planning
**Scenario:** Plan synthesis route for target molecule.
**Steps:**
1. Load dataset: `datasets.USPTO50k()`
2. Train center identification model (RGCN)
3. Train synthon completion model (GIN)
4. Combine into end-to-end retrosynthesis pipeline
5. Apply recursively for multi-step planning
6. Check commercial availability of building blocks
**Navigation:** `references/retrosynthesis.md` → Task types → Multi-step planning
## Integration Patterns
### With RDKit
Convert between TorchDrug molecules and RDKit:
```python
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:
```python
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:
```python
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:
- Architecture philosophy (modular, configurable)
- Data structures (Graph, Molecule, Protein, PackedGraph)
- Model interface and forward function signature
- Task interface (predict, target, forward, evaluate)
- Training workflows and best practices
- Loss functions and metrics
- Common pitfalls and debugging
## Quick Reference Cheat Sheet
**Choose Dataset:**
- Molecular property → `references/datasets.md` → Molecular section
- Protein task → `references/datasets.md` → Protein section
- Knowledge graph → `references/datasets.md` → Knowledge graph section
**Choose Model:**
- Molecules → `references/models_architectures.md` → GNN section → GIN/GAT/SchNet
- Proteins (sequence) → `references/models_architectures.md` → Protein section → ESM
- Proteins (structure) → `references/models_architectures.md` → Protein section → GearNet
- Knowledge graph → `references/models_architectures.md` → KG section → RotatE/ComplEx
**Common Tasks:**
- Property prediction → `references/molecular_property_prediction.md` or `references/protein_modeling.md`
- Generation → `references/molecular_generation.md`
- Retrosynthesis → `references/retrosynthesis.md`
- KG reasoning → `references/knowledge_graphs.md`
**Understand Architecture:**
- Data structures → `references/core_concepts.md` → Data Structures
- Model design → `references/core_concepts.md` → Model Interface
- Task design → `references/core_concepts.md` → Task Interface
## 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
## Resources
**Official Documentation:** https://torchdrug.ai/docs/
**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:
1. **Molecular property prediction** → `molecular_property_prediction.md`
2. **Protein modeling** → `protein_modeling.md`
3. **Knowledge graphs** → `knowledge_graphs.md`
4. **Molecular generation** → `molecular_generation.md`
5. **Retrosynthesis** → `retrosynthesis.md`
6. **Model selection** → `models_architectures.md`
7. **Dataset selection** → `datasets.md`
8. **Technical details** → `core_concepts.md`
Each reference provides comprehensive coverage of its domain with examples, best practices, and common use cases.