ContextQMD
Back to Claude Scientific Skills

Molecular Property Prediction and Pretraining

skills/torchdrug/references/molecular_property_prediction.md

2.54.05.9 KB
Original Source

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:

python
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

python
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

python
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

python
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:

python
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:

python
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:

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

InfoGraph

python
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

python
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:

python
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.