Molecular Property Prediction and Pretraining

Follow the official property prediction and pretrained molecular representations tutorials for TorchDrug 0.2.1.

Supervised property prediction

1. Load and split data

The official tutorial uses a random 80/10/10 ClinTox split:

import torch
from torchdrug import datasets

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths.append(len(dataset) - sum(lengths))
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)

This is a random split, not a scaffold split. If a benchmark requires a scaffold split, implement or import that protocol explicitly and record it in the experiment configuration.

2. Define the representation model

from torchdrug import models

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)

3. Define the task

from torchdrug import tasks

task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=("auprc", "auroc"),
)

task means the target field name(s) or a mapping of target names to weights. It does not mean "node", "edge", or "graph".

Documented PropertyPrediction criteria are:

"mse"
"bce"
"ce"

Documented metrics are:

"mae"
"rmse"
"auprc"
"auroc"

Other useful constructor options include num_mlp_layer, normalization, num_class, mlp_batch_norm, mlp_dropout, and graph_construction_model.

For large multi-label problems, inspect tasks.MultipleBinaryClassification, which has its own task IDs, metrics, and reweighting behavior.

4. Train with Engine

from torchdrug import core

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=1024,
)
solver.train(num_epoch=100)
solver.evaluate("valid")

Add gpus=[0] only for supported CUDA execution. Start with one epoch and a smaller batch for a smoke test.

Manual prediction

Use TorchDrug collation:

from torch.nn import functional as F
from torchdrug import data

batch = data.graph_collate(valid_set[:8])
logits = task.predict(batch)
probabilities = F.sigmoid(logits)
targets = task.target(batch)

For binary classification, predict() returns logits and the tutorial applies sigmoid. For normalized regression, TorchDrug 0.2.1 returns predictions on the original target scale; this changed from earlier releases.

When predicting on CUDA manually, move the whole nested batch:

from torchdrug import utils

batch = utils.cuda(batch)

Self-supervised pretraining

The tutorial uses ClinTox only as a small illustration and recommends a larger unlabeled corpus such as ZINC2m for real pretraining.

Use matching pretraining features:

dataset = datasets.ClinTox(
    "~/molecule-datasets/",
    atom_feature="pretrain",
    bond_feature="pretrain",
)

InfoGraph

from torchdrug import core, models, tasks

gin_model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[300, 300, 300, 300, 300],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean",
)
model = models.InfoGraph(gin_model, separate_model=False)
task = tasks.Unsupervised(model)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    dataset,
    None,
    None,
    optimizer,
    batch_size=256,
)
solver.train(num_epoch=100)
solver.save("gin-infograph.pth")

Attribute masking

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[300, 300, 300, 300, 300],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean",
)
task = tasks.AttributeMasking(model, mask_rate=0.15)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    dataset,
    None,
    None,
    optimizer,
    batch_size=256,
)
solver.train(num_epoch=100)
solver.save("gin-attribute-masking.pth")

Fine-tune the encoder

Recreate the same GIN architecture and feature dimensions, then wrap it in the supervised task:

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[300, 300, 300, 300, 300],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean",
)
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=("auprc", "auroc"),
)

checkpoint = torch.load("gin-attribute-masking.pth")["model"]
task.load_state_dict(checkpoint, strict=False)

Then construct a new optimizer and supervised Engine. strict=False is intentional because the pretraining and supervised task heads differ. Review missing and unexpected keys if changing the architecture.

Experiment checks

Confirm dataset.tasks names and label shapes.
Confirm classification vs regression before choosing criterion and metrics.
Record the exact split protocol; do not mislabel random splits as scaffold splits.
Use AUPRC as well as AUROC for heavily imbalanced binary tasks.
Keep feature arguments identical when loading pretrained weights.
Fit preprocessing only on the training split.
Reserve the test split until model selection is complete.