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:

• 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