Molecules as tensors: a PyTorchic take on drug discovery

What it does TorchDrug is a PyTorch-based toolbox for graph machine learning with a specialization in drug discovery. It provides Graph and Molecule objects that support GPU acceleration, auto differentiation, and subgraph extraction, along with off-the-shelf datasets like Tox21 and models such as GIN for property prediction. The library aims to lower the barrier for ML practitioners who want to experiment on molecular data but lack domain expertise in chemistry.

The interesting bit The API treats molecular structures as first-class tensor citizens: you can instantiate a molecule from a SMILES string, register custom bond or atom attributes, and index into subgraphs while the framework automatically tracks those attributes. Scaling from a single CPU to distributed multi-GPU training is handled by the same core.Engine class with only a device-list argument changing.

Key highlights

• Native PyTorch graph operations with GPU and autograd support
• Molecule objects parsed from SMILES strings, including scaffold extraction and custom node/edge attributes
• Built-in datasets and standard models (e.g., GIN) for rapid prototyping
• Single-line switching between CPU, single-GPU, multi-GPU, and distributed training via core.Engine
• Built-in experiment tracking integration for Weights & Biases

Caveats

• Installation can be involved: Windows requires Visual Studio build tools and PowerShell environment setup, while Apple Silicon lacks MPS support and needs torch-scatter and torch-cluster compiled from source
• Python support is capped at version 3.10, and PyTorch compatibility starts at 1.8.0 with some install docs pinning older point releases

Verdict Worth a look if you are a machine-learning researcher or practitioner who wants to run graph neural networks on molecular data without writing cheminformatics glue code. Pure graph-learning specialists outside chemistry may find the drug-discovery scaffolding unnecessary.