Core Concepts and Technical Details

Overview

This reference covers TorchDrug's fundamental architecture, design principles, and technical implementation details.

Architecture Philosophy

Modular Design

TorchDrug separates concerns into distinct modules:

1. Representation Models (models.py): Encode graphs into embeddings
2. Task Definitions (tasks.py): Define learning objectives and evaluation
3. Data Handling (data.py, datasets.py): Graph structures and datasets
4. Core Components (core.py): Base classes and utilities

Benefits: - Reuse representations across tasks - Mix and match components - Easy experimentation and prototyping - Clear separation of concerns

Configurable System

All components inherit from core.Configurable: - Serialize to configuration dictionaries - Reconstruct from configurations - Save and load complete pipelines - Reproducible experiments

Core Components

core.Configurable

Base class for all TorchDrug components.

Key Methods: - config_dict(): Serialize to dictionary - load_config_dict(config): Load from dictionary - save(file): Save to file - load(file): Load from file

Example:

from torchdrug import core, models

model = models.GIN(input_dim=10, hidden_dims=[256, 256])

# Save configuration
config = model.config_dict()
# {'class': 'GIN', 'input_dim': 10, 'hidden_dims': [256, 256], ...}

# Reconstruct model
model2 = core.Configurable.load_config_dict(config)

core.Registry

Decorator for registering models, tasks, and datasets.

Usage:

from torchdrug import core as core_td

@core_td.register("models.CustomModel")
class CustomModel(nn.Module, core_td.Configurable):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, hidden_dim)

    def forward(self, graph, input, all_loss, metric):
        # Model implementation
        pass

Benefits: - Models automatically serializable - String-based model specification - Easy model lookup and instantiation

Data Structures

Graph

Core data structure representing molecular or protein graphs.

Attributes: - num_node: Number of nodes - num_edge: Number of edges - node_feature: Node feature tensor [num_node, feature_dim] - edge_feature: Edge feature tensor [num_edge, feature_dim] - edge_list: Edge connectivity [num_edge, 2 or 3] - num_relation: Number of edge types (for multi-relational)

Methods: - node_mask(mask): Select subset of nodes - edge_mask(mask): Select subset of edges - undirected(): Make graph undirected - directed(): Make graph directed

Batching: - Graphs batched into single disconnected graph - Automatic batching in DataLoader - Preserves node/edge indices per graph

Molecule (extends Graph)

Specialized graph for molecules.

Additional Attributes: - atom_type: Atomic numbers - bond_type: Bond types (single, double, triple, aromatic) - formal_charge: Atomic formal charges - explicit_hs: Explicit hydrogen counts

Methods: - from_smiles(smiles): Create from SMILES string - from_molecule(mol): Create from RDKit molecule - to_smiles(): Convert to SMILES - to_molecule(): Convert to RDKit molecule - ion_to_molecule(): Neutralize charges

Example:

from torchdrug import data

# From SMILES
mol = data.Molecule.from_smiles("CCO")

# Atom features
print(mol.atom_type)  # [6, 6, 8] (C, C, O)
print(mol.bond_type)  # [1, 1] (single bonds)

Protein (extends Graph)

Specialized graph for proteins.

Additional Attributes: - residue_type: Amino acid types - atom_name: Atom names (CA, CB, etc.) - atom_type: Atomic numbers - residue_number: Residue numbering - chain_id: Chain identifiers

Methods: - from_pdb(pdb_file): Load from PDB file - from_sequence(sequence): Create from sequence - to_pdb(pdb_file): Save to PDB file

Graph Construction: - Nodes typically represent residues (not atoms) - Edges can be sequential, spatial (KNN), or contact-based - Configurable edge construction strategies

Example:

from torchdrug import data

# Load protein
protein = data.Protein.from_pdb("1a3x.pdb")

# Build graph with multiple edge types
graph = protein.residue_graph(
    node_position="ca",  # Use Cα positions
    edge_types=["sequential", "radius"]  # Sequential + spatial edges
)

PackedGraph

Efficient batching structure for heterogeneous graphs.

Purpose: - Batch graphs of different sizes - Single GPU memory allocation - Efficient parallel processing

Attributes: - num_nodes: List of node counts per graph - num_edges: List of edge counts per graph - graph_ind: Graph index for each node

Use Cases: - Automatic in DataLoader - Custom batching strategies - Multi-graph operations

Model Interface

Forward Function Signature

All TorchDrug models follow a standardized interface:

def forward(self, graph, input, all_loss=None, metric=None):
    """
    Args:
        graph (Graph): Batch of graphs
        input (Tensor): Node input features
        all_loss (Tensor, optional): Accumulator for losses
        metric (dict, optional): Dictionary for metrics

    Returns:
        dict: Output dictionary with representation keys
    """
    # Model computation
    output = self.layers(graph, input)

    return {
        "node_feature": output,
        "graph_feature": graph_pooling(output)
    }

Key Points: - graph: Batched graph structure - input: Node features [num_node, input_dim] - all_loss: Accumulated loss (for multi-task) - metric: Shared metric dictionary - Returns dict with representation types

Essential Attributes

All models must define: - input_dim: Expected input feature dimension - output_dim: Output representation dimension

Purpose: - Automatic dimension checking - Compose models in pipelines - Error checking and validation

Example:

class CustomModel(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = hidden_dim
        # ... layers ...

Task Interface

Core Task Methods

All tasks implement these methods:

class CustomTask(tasks.Task):
    def preprocess(self, train_set, valid_set, test_set):
        """Dataset-specific preprocessing (optional)"""
        pass

    def predict(self, batch):
        """Generate predictions for a batch"""
        graph, label = batch
        output = self.model(graph, graph.node_feature)
        pred = self.mlp(output["graph_feature"])
        return pred

    def target(self, batch):
        """Extract ground truth labels"""
        graph, label = batch
        return label

    def forward(self, batch):
        """Compute training loss"""
        pred = self.predict(batch)
        target = self.target(batch)
        loss = self.criterion(pred, target)
        return loss

    def evaluate(self, pred, target):
        """Compute evaluation metrics"""
        metrics = {}
        metrics["auroc"] = compute_auroc(pred, target)
        metrics["auprc"] = compute_auprc(pred, target)
        return metrics

Task Components

Typical Task Structure: 1. Representation Model: Encodes graph to embeddings 2. Readout/Prediction Head: Maps embeddings to predictions 3. Loss Function: Training objective 4. Metrics: Evaluation measures

Example:

from torchdrug import tasks, models

# Representation model
model = models.GIN(input_dim=10, hidden_dims=[256, 256])

# Task wraps model with prediction head
task = tasks.PropertyPrediction(
    model=model,
    task=["task1", "task2"],  # Multi-task
    criterion="bce",
    metric=["auroc", "auprc"],
    num_mlp_layer=2
)

Training Workflow

Standard Training Loop

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

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

# 2. Create data loaders
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=32)

# 3. Define model and task
model = models.GIN(input_dim=dataset.node_feature_dim,
                   hidden_dims=[256, 256, 256])
task = tasks.PropertyPrediction(model, task=dataset.tasks,
                                 criterion="bce", metric=["auroc", "auprc"])

# 4. Setup optimizer
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)

# 5. Training loop
for epoch in range(100):
    # Train
    task.train()
    for batch in train_loader:
        loss = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Validate
    task.eval()
    preds, targets = [], []
    for batch in valid_loader:
        pred = task.predict(batch)
        target = task.target(batch)
        preds.append(pred)
        targets.append(target)

    preds = torch.cat(preds)
    targets = torch.cat(targets)
    metrics = task.evaluate(preds, targets)
    print(f"Epoch {epoch}: {metrics}")

PyTorch Lightning Integration

TorchDrug tasks are compatible with PyTorch Lightning:

import pytorch_lightning as pl

class LightningWrapper(pl.LightningModule):
    def __init__(self, task):
        super().__init__()
        self.task = task

    def training_step(self, batch, batch_idx):
        loss = self.task(batch)
        return loss

    def validation_step(self, batch, batch_idx):
        pred = self.task.predict(batch)
        target = self.task.target(batch)
        return {"pred": pred, "target": target}

    def validation_epoch_end(self, outputs):
        preds = torch.cat([o["pred"] for o in outputs])
        targets = torch.cat([o["target"] for o in outputs])
        metrics = self.task.evaluate(preds, targets)
        self.log_dict(metrics)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

Loss Functions

Built-in Criteria

Classification: - "bce": Binary cross-entropy - "ce": Cross-entropy (multi-class)

Regression: - "mse": Mean squared error - "mae": Mean absolute error

Knowledge Graph: - "bce": Binary classification of triples - "ce": Cross-entropy ranking loss - "margin": Margin-based ranking

Custom Loss

class CustomTask(tasks.Task):
    def forward(self, batch):
        pred = self.predict(batch)
        target = self.target(batch)

        # Custom loss computation
        loss = custom_loss_function(pred, target)

        return loss

Metrics

Common Metrics

Classification: - AUROC: Area under ROC curve - AUPRC: Area under precision-recall curve - Accuracy: Overall accuracy - F1: Harmonic mean of precision and recall

Regression: - MAE: Mean absolute error - RMSE: Root mean squared error - R²: Coefficient of determination - Pearson: Pearson correlation

Ranking (Knowledge Graph): - MR: Mean rank - MRR: Mean reciprocal rank - Hits@K: Percentage in top K

Multi-Task Metrics

For multi-label or multi-task: - Metrics computed per task - Macro-average across tasks - Can weight by task importance

Data Transforms

Molecule Transforms

from torchdrug import transforms

# Add virtual node connected to all atoms
transform1 = transforms.VirtualNode()

# Add virtual edges
transform2 = transforms.VirtualEdge()

# Compose transforms
transform = transforms.Compose([transform1, transform2])

dataset = datasets.BBBP("~/datasets/", transform=transform)

Protein Transforms

# Add edges based on spatial proximity
transform = transforms.TruncateProtein(max_length=500)

dataset = datasets.Fold("~/datasets/", transform=transform)

Best Practices

Memory Efficiency

1. Gradient Accumulation: For large models
2. Mixed Precision: FP16 training
3. Batch Size Tuning: Balance speed and memory
4. Data Loading: Multiple workers for I/O

Reproducibility

1. Set Seeds: PyTorch, NumPy, Python random
2. Deterministic Operations: torch.use_deterministic_algorithms(True)
3. Save Configurations: Use core.Configurable
4. Version Control: Track TorchDrug version

Debugging

1. Check Dimensions: Verify input_dim and output_dim
2. Validate Batching: Print batch statistics
3. Monitor Gradients: Watch for vanishing/exploding
4. Overfit Small Batch: Ensure model capacity

Performance Optimization

1. GPU Utilization: Monitor with nvidia-smi
2. Profile Code: Use PyTorch profiler
3. Optimize Data Loading: Prefetch, pin memory
4. Compile Models: Use TorchScript if possible

Advanced Topics

Multi-Task Learning

Train single model on multiple related tasks:

task = tasks.PropertyPrediction(
    model,
    task=["task1", "task2", "task3"],
    criterion="bce",
    metric=["auroc"],
    task_weight=[1.0, 1.0, 2.0]  # Weight task 3 more
)

Transfer Learning

1. Pre-train on large dataset
2. Fine-tune on target dataset
3. Optionally freeze early layers

Self-Supervised Pre-training

Use pre-training tasks: - AttributeMasking: Mask node features - EdgePrediction: Predict edge existence - ContextPrediction: Contrastive learning

Custom Layers

Extend TorchDrug with custom GNN layers:

from torchdrug import layers

class CustomConv(layers.MessagePassingBase):
    def message(self, graph, input):
        # Custom message function
        pass

    def aggregate(self, graph, message):
        # Custom aggregation
        pass

    def combine(self, input, update):
        # Custom combination
        pass

Common Pitfalls

1. Forgetting input_dim and output_dim: Models won't compose
2. Not Batching Properly: Use PackedGraph for variable-sized graphs
3. Data Leakage: Be careful with scaffold splits and pre-training
4. Ignoring Edge Features: Bonds/spatial info can be critical
5. Wrong Evaluation Metrics: Match metrics to task (AUROC for imbalanced)
6. Insufficient Regularization: Use dropout, weight decay, early stopping
7. Not Validating Chemistry: Generated molecules must be valid
8. Overfitting Small Datasets: Use pre-training or simpler models