Part 4: Train a Model#

Table of Contents#

Introduction#

This is the fourth and final tutorial in the How to Build a Model Using AtomWorks series. By now you have:

  • Parquet files for your train/validation/test splits (Part 1).

  • A loader that turns each parquet row into a protein-ligand structure (Part 2).

  • A transform pipeline that crops the pocket and converts it into tensors (Part 2).

  • A PocketDockGNN model in model.py (Part 3).

In this installment, you will write train.py, which wires all of these pieces together to train, validate, checkpoint, and test the model.

Important

This tutorial completes the training script using the AtomWorks API, PyTorch, and PyTorch Lightning. The full solution is available in the tutorial files, and the code below is hidden in collapsible cells so you can attempt each step yourself first.

Prerequisites#

Before starting this part it is assumed that you have:

  • Completed Parts 1–3, with transforms.py and model.py in place.

  • A working installation of PyTorch and PyTorch Lightning.

  • A CPU or GPU; the trainer chooses an available accelerator.

What the Training Script Does#

train.py wires together four pieces:

  1. Read the split parquets.

  2. Load and transform each example into model inputs.

  3. Instantiate your model (PocketDockGNN).

  4. Train, validate, checkpoint, and test.

The model sees a pocket-centered graph with five tensors:

  • atomic_numbers: atomic identity for each atom.

  • input_coords: protein pocket coordinates are kept, ligand coordinates are zeroed out.

  • target_coords: the true coordinates the model should predict.

  • edge_index: bond graph connectivity.

  • is_ligand: a boolean mask for which atoms belong to the ligand; the model uses its complement to pool protein pocket coordinates.

In other words, the task is: given the protein pocket context and the ligand atoms, predict the ligand’s 3D placement.

Imports and Global Settings#

In a new file, train.py, import PyTorch, PyTorch Lightning, the AtomWorks dataset utilities used in the previous scripts, the transforms you wrote in Part 2, and the model class from Part 3.

Click to see the code
import torch
import pandas as pd
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from torch.utils.data import DataLoader, Dataset
from pathlib import Path
import argparse

from atomworks.ml.datasets import PandasDataset
from atomworks.ml.datasets.loaders import create_structure_loader
from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms
from atomworks.ml.transforms.base import Compose, ConvertToTorch

from transforms import CropToPocket, FeaturizeForDocking
from model import PocketDockGNN

Note that we now also import ConvertToTorch, which turns the NumPy features from FeaturizeForDocking into torch tensors.

Add two small global settings:

  • torch.set_float32_matmul_precision("medium"): a practical speed/precision trade-off for training.

  • pl.seed_everything(42): makes runs more reproducible.

Click to see the code
torch.set_float32_matmul_precision("medium")
pl.seed_everything(42)

# Call these in main() before training.

Configuration#

Put the tunable settings from the previous scripts and the maximum sizes for our training, validation, and testing datasets in a CONFIG dictionary near the top of the file.

Click to see the code
CONFIG = {
    "hidden_dim": 128,
    "num_layers": 3,
    "learning_rate": 1e-3,
    "max_epochs": 5,
    "pocket_radius": 10.0,
    "num_workers": 0,
    "max_train": 100,
    "max_val": 20,
    "max_test": 20,
}

Note

The max_* values cap how many examples we use so the run stays small and fast to troubleshoot. For actual training, remove these caps (set them to None) and train on the full splits.

List the Tensor Keys#

Your featurization step produces five tensors. Store those once in a TENSOR_KEYS list so the rest of the file can reuse the same list. This list tells ConvertToTorch what to convert and the collate function what to stack.

Click to see the code
TENSOR_KEYS = [
    "atomic_numbers",
    "input_coords",
    "target_coords",
    "edge_index",
    "is_ligand",
]

Important

TENSOR_KEYS must be defined before the functions that reference it (collate_fn and build_pipeline). Keep it near the top of the file with CONFIG.

Make the Dataset Robust to Bad Examples#

If you train on real structures (which we are) some examples will fail. A ligand may have no resolved coordinates, or the pocket crop may remove everything useful. You do not want one bad structure to crash the entire run.

Wrap PandasDataset in a small RobustDataset that inherits from torch.utils.data.Dataset. If a single example raises during loading or transform, it records the index and returns None instead of killing training.

Click to see the code
class RobustDataset(Dataset):
    def __init__(self, dataset: PandasDataset):
        self.dataset = dataset
        self.failed = []

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        try:
            return self.dataset[idx]
        except (OSError, ValueError, KeyError, IndexError) as exc:
            print(f"Skipping example {idx}: {exc}")
            self.failed.append(idx)
            return None

Write a Custom collate_fn#

The default PyTorch collator is not a good fit here: the examples are variable-size graphs, the dataset may return None, and we only want to batch the specific tensors the model needs (ignoring any extra fields in the example dictionary). Your collate function should do two things: drop failed examples, and stack only the tensor keys the model expects.

Click to see the code
def collate_fn(batch):
    batch = [b for b in batch if b is not None]
    if len(batch) == 0:
        return None

    return {
        k: torch.stack([example[k] for example in batch])
        for k in TENSOR_KEYS
    }

With batch_size=1, this stacks a single example and adds a leading dimension of size 1, which the model’s _shared_step removes with .squeeze(0).

Rebuild the Transform Pipeline#

Write a helper that reconstructs the same transform sequence you tested in Part 2, now with ConvertToTorch appended so the features come out as torch tensors.

Click to see the code
def build_pipeline(radius: float) -> Compose:
    return Compose([
        RemoveHydrogens(),
        RemoveUnresolvedAtoms(),
        CropToPocket(radius=radius),
        FeaturizeForDocking(),
        ConvertToTorch(keys=TENSOR_KEYS),
    ])

Dataset and Dataloader Factories#

Write one function that takes a split parquet path and returns a ready-to-use dataset (this is essentially what you did in smoke_test.py, wrapped in RobustDataset).

Click to see the code
def build_dataset(parquet_path: str, name: str, radius: float, max_examples: int | None = None) -> RobustDataset:
    df = pd.read_parquet(parquet_path)
    if max_examples is not None:
        df = df.head(max_examples).reset_index(drop=True)

    loader = create_structure_loader(
        altloc_seed_colname="altloc_seed",
        column_mapping={
            "query_pn_unit_iids": ["pn_unit_1_iid", "pn_unit_2_iid"],
            "query_is_polymer": ["pn_unit_1_is_polymer", "pn_unit_2_is_polymer"],
        },
    )

    dataset = PandasDataset(
        data=df,
        name=name,
        id_column="example_id",
        loader=loader,
        transform=build_pipeline(radius),
    )
    return RobustDataset(dataset)

The wrapper prints the failing row index and exception so skipped examples remain visible. Keep at least one valid example in each split.

Then write one small helper to build dataloaders consistently. This keeps the script tidy and reuses the same settings for train, validation, and test.

Click to see the code
def build_dataloader(dataset, shuffle: bool) -> DataLoader:
    return DataLoader(
        dataset,
        batch_size=1,
        shuffle=shuffle,
        num_workers=CONFIG["num_workers"],
        collate_fn=collate_fn,
        persistent_workers=False,
    )

Build the Datasets and Dataloaders#

Now the main script body. Start by building the three datasets. The downloadable script accepts --splits-dir, --output-dir, --max-epochs, and example caps. Set output_dir = Path("tutorial_run") for the snippets below.

Click to see the code
print("Building datasets...")
train_dataset = build_dataset(
    "splits/train.parquet", "docking_train", CONFIG["pocket_radius"], CONFIG["max_train"]
)
val_dataset = build_dataset(
    "splits/validation.parquet", "docking_val", CONFIG["pocket_radius"], CONFIG["max_val"]
)
test_dataset = build_dataset(
    "splits/test.parquet", "docking_test", CONFIG["pocket_radius"], CONFIG["max_test"]
)

print(f"  Train: {len(train_dataset):,} examples")
print(f"  Val:   {len(val_dataset):,} examples")
print(f"  Test:  {len(test_dataset):,} examples")

Then build the dataloaders.

Click to see the code
train_loader = build_dataloader(train_dataset, shuffle=True)
val_loader = build_dataloader(val_dataset, shuffle=False)
test_loader = build_dataloader(test_dataset, shuffle=False)

Instantiate the Model#

Instantiate the network with the hyperparameters from CONFIG.

Click to see the code
model = PocketDockGNN(
    hidden_dim=CONFIG["hidden_dim"],
    num_layers=CONFIG["num_layers"],
    learning_rate=CONFIG["learning_rate"],
)

Optionally, print the parameter count as a sanity check that you built the expected network:

print(f"\nModel parameters: {sum(p.numel() for p in model.parameters()):,}")

Checkpointing and Early Stopping#

Before creating the trainer, define the callbacks. Put checkpoints under output_dir. Use periodic_checkpoint to give recovery points during training, and best_checkpoint to keep the best models by val/loss.

Click to see the code
periodic_checkpoint = ModelCheckpoint(
    dirpath=output_dir / "checkpoints",
    filename="pocketdockgnn-{epoch:02d}-{step}",
    every_n_train_steps=50,
    save_last=True,
    verbose=True,
)

best_checkpoint = ModelCheckpoint(
    dirpath=output_dir / "checkpoints",
    filename="pocketdockgnn-best-{epoch:02d}",
    monitor="val/loss",
    mode="min",
    save_top_k=3,
    verbose=True,
)

callbacks = [
    best_checkpoint,
    periodic_checkpoint,
    EarlyStopping(
        monitor="val/loss",
        patience=10,
        mode="min",
        verbose=True,
    ),
]

The EarlyStopping callback stops training if validation loss stops improving for patience checks.

Create the Trainer#

Now define the trainer itself with the PyTorch Lightning Trainer.

Click to see the code
trainer = pl.Trainer(
    max_epochs=CONFIG["max_epochs"],
    accelerator="auto",
    devices=1,
    callbacks=callbacks,
    log_every_n_steps=1,
    default_root_dir=output_dir,
    check_val_every_n_epoch=1,
    enable_progress_bar=True,
)

Tip

The small capped example runs on CPU when no GPU is available.

Train and Evaluate#

With train.parquet, validation.parquet, and test.parquet in splits/, run the downloaded scripts from one directory:

python train.py --splits-dir splits --output-dir tutorial_run

Call fit() to train.

Click to see the code
print("\nStarting training...")
trainer.fit(
    model,
    train_dataloaders=train_loader,
    val_dataloaders=val_loader,
)

After training, print some diagnostics: how many examples were skipped, and which checkpoint Lightning considered best.

Click to see the code
print(f"\nFailed examples during training: {len(train_dataset.failed)}")
print(f"Failed examples during val:      {len(val_dataset.failed)}")
print(f"Best checkpoint:                {best_checkpoint.best_model_path}")
if not best_checkpoint.best_model_path:
    raise RuntimeError("No validation checkpoint was saved; inspect skipped examples")

Finally, evaluate on the held-out test split using the best checkpoint.

Click to see the code
print("\nEvaluating on test set...")
trainer.test(
    model,
    dataloaders=test_loader,
    ckpt_path=best_checkpoint.best_model_path,
)

print("\nDone.")

Wrapping Up#

You have now built a complete, if simplified, machine-learning pipeline with AtomWorks: from raw PDB metadata to a trained pose-generation model. As next steps you might:

  • Remove the max_* caps and train on the full splits.

  • Add symmetric edges or an equivariant architecture to remove the model’s dependence on absolute coordinates (see the warning in Part 3).

  • Track additional metrics, such as per-example ligand RMSD distributions, or visualize predicted poses.

Running Inference on a Trained Model#

trainer.test() reports metrics, but it does not hand you the predicted poses. To get coordinates for a new pocket-ligand example, reload the best checkpoint and call the model directly. You can find an inference.py script in the tutorial files. After training, pass the best checkpoint path printed by train.py:

python inference.py PATH_TO_BEST_CHECKPOINT --split splits/test.parquet
Click to see the code
model = PocketDockGNN.load_from_checkpoint(best_checkpoint.best_model_path, map_location="cpu")
model.eval()

batch = next((example for example in test_loader if example is not None), None)
if batch is None:
    raise ValueError("No usable examples in inference split")
atomic_numbers = batch["atomic_numbers"].squeeze(0)
input_coords = batch["input_coords"].squeeze(0)
edge_index = batch["edge_index"].squeeze(0)
is_ligand = batch["is_ligand"].squeeze(0)

with torch.no_grad():
    pred_coords = model(atomic_numbers, input_coords, edge_index, is_ligand)

predicted_ligand_coords = pred_coords[is_ligand]

The .squeeze(0) calls undo the batch dimension DataLoader adds, matching what _shared_step does internally. Only the rows where is_ligand is True are ligand predictions. Other rows are unconstrained protein predictions and should not be interpreted as a fixed pocket.

Note

pred_coords is a raw coordinate prediction, not validated molecular geometry. Check bond lengths and clashes before treating it as a usable pose.

Glossary#

Collate function: a function that assembles a list of examples into a single batch; here it also drops failed (None) examples.

Checkpoint: a saved snapshot of model weights (and optimizer state) that can be reloaded to resume training or run evaluation.

Early stopping: halting training once a monitored metric (here val/loss) stops improving, to save compute and reduce overfitting.