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/val/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.

  • Access to a GPU (the example trainer is configured for a single GPU, but you can change the 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.

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, model.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 atomworks.ml.datasets import PandasDataset
from atomworks.ml.datasets.loaders import create_loader_with_query_pn_units
from atomworks.ml.transforms.filters import RemoveHydrogens, RemoveUnresolvedAtoms
from atomworks.ml.transforms.base import ConvertToTorch, Compose

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)

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,
    "batch_size": 1,
    "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 is used both to tell ConvertToTorch what to convert and to tell 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 Exception:
            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
        if k in batch[0]
    }

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):
    df = pd.read_parquet(parquet_path)
    if max_examples is not None:
        df = df.head(max_examples).reset_index(drop=True)

    loader = create_loader_with_query_pn_units(
        pn_unit_iid_colnames=["pn_unit_1_iid", "pn_unit_2_iid"]
    )

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

save_failed_examples_to_dir writes any example that fails inside the AtomWorks pipeline to disk so you can inspect it later.

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=CONFIG["batch_size"],
        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.

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/val.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. 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="checkpoints/",
    filename="pocketdockgnn-{epoch:02d}-{step}",
    every_n_train_steps=50,
    save_last=True,
    verbose=True,
)

best_checkpoint = ModelCheckpoint(
    dirpath="checkpoints/",
    filename="pocketdockgnn-best-{epoch:02d}-{val/loss:.4f}",
    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="gpu",
    devices=1,
    callbacks=callbacks,
    log_every_n_steps=10,
    val_check_interval=50,
    enable_progress_bar=True,
)

Tip

If you do not have a GPU available, set accelerator="cpu" (and drop devices=1) to run the small capped example on CPU.

Train and Evaluate#

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}")

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 a inference.py script in the tutorial files.

Click to see the code
model = PocketDockGNN.load_from_checkpoint(best_checkpoint.best_model_path)
model.eval()

batch = next(iter(test_loader))
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)

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 the coordinates you asked the model to predict. The rest are the (unchanged) protein pocket context.

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.