Tune PyTorch models with Polyaxon

A PyTorch training loop can stay the same while you try different model sizes, optimizers, and learning rates. Polyaxon can choose those configurations, run each one on a GPU, and use the validation results to guide the next batch.

This walkthrough trains a small MNIST classifier and tunes five parameters with native TPE, using Polyaxon's search matrix and tracked outputs to manage the study.

A search space feeds configurations into GPU training trials. Validation accuracy returns to native TPE to guide the next batch.

Define what to tune

We will choose the number of hidden layers, their shared width, dropout probability, optimizer, and learning rate. Every hidden layer uses the same width, so all five parameters affect every trial.

Use a deployment and CLI with Polyaxon 2.17's native TPE support, a configured project and artifact store, and an NVIDIA GPU queue. Native TPE uses the built-in tpe-tuner component to generate suggestions; your deployment must be able to resolve and run it. The tuner configuration lets you choose its queue or use a custom component. Containers need access to the image registry, Python packages, and MNIST downloads.

Create a folder containing the three files below. Start with study.yaml:

version: 1.1
kind: operation
name: mnist-tpe
pathRef: ./train-component.yaml
queue: research/gpu
params:
  epochs: 5
  seed: 23
matrix:
  kind: tpe
  numRuns: 10
  maxIterations: 2
  concurrency: 2
  seed: 19
  metric:
    name: validation_accuracy
    optimization: maximize
  params:
    n_layers:
      kind: choice
      value: [1, 2, 3]
    hidden_width:
      kind: choice
      value: [32, 64, 128, 256]
    dropout:
      kind: uniform
      value: [0.0, 0.5]
    optimizer:
      kind: choice
      value: [Adam, RMSprop, SGD]
    learning_rate:
      kind: loguniform
      value: [-9.21034, -2.30259]  # ln(0.0001), ln(0.1)

Polyaxon's loguniform bounds are in natural-log space: the learning rate is sampled between approximately 0.0001 and 0.1. Dropout is sampled from [0.0, 0.5). See the distribution reference when changing the ranges.

numRuns: 10 generates ten suggestions per batch. maxIterations: 2 allows the initial batch at iteration 0 followed by iterations 1 and 2: up to 30 training trials. concurrency: 2 allows two trial operations at once, subject to queue capacity.

The first batch samples randomly. Later batches use successful trials with a reported objective to guide suggestions as observations accumulate. Each batch completes before the next is generated; this is a batch study rather than a new suggestion after every individual completion. The native TPE walkthrough explains the sampler in more detail.

Give each trial a GPU

Save this as train-component.yaml:

version: 1.1
kind: component
name: mnist-train
termination:
  timeout: 900
  maxRetries: 0
inputs:
  - name: n_layers
    type: int
  - name: hidden_width
    type: int
  - name: dropout
    type: float
  - name: optimizer
    type: str
  - name: learning_rate
    type: float
  - name: epochs
    type: int
    value: 5
  - name: batch_size
    type: int
    value: 128
  - name: seed
    type: int
    value: 23
outputs:
  - name: validation_accuracy
    type: float
run:
  kind: job
  container:
    image: pytorch/pytorch:2.7.1-cuda12.6-cudnn9-runtime
    workingDir: "{{ globals.run_artifacts_path }}/uploads"
    command: [bash, -c]
    args:
      - |
        set -euo pipefail
        python -m pip install --no-cache-dir polyaxon torchvision==0.22.1
        python -m pip freeze
        exec python train.py \
          --n-layers {{ n_layers }} \
          --hidden-width {{ hidden_width }} \
          --dropout {{ dropout }} \
          --optimizer {{ optimizer }} \
          --learning-rate {{ learning_rate }} \
          --epochs {{ epochs }} \
          --batch-size {{ batch_size }} \
          --seed {{ seed }}
    resources:
      requests:
        cpu: "2"
        memory: "4Gi"
        nvidia.com/gpu: 1
      limits:
        cpu: "4"
        memory: "8Gi"
        nvidia.com/gpu: 1

Replace research/gpu in the study with your allowed agent/queue. Each trial requests one GPU on that agent's Kubernetes cluster, which needs compatible NVIDIA drivers and device-plugin support. With two trials running, training can occupy two GPUs. Each trial has a 15-minute timeout and no retries.

The PyTorch image supplies PyTorch 2.7.1 and CUDA 12.6; torchvision 0.22.1 is its matching release. The tracking client is installed at runtime. The package list goes into the logs; pin your reviewed client version and image digest when repeating the study.

Train and report validation accuracy

Save this as train.py:

import argparse
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

from polyaxon import tracking

parser = argparse.ArgumentParser()
parser.add_argument("--n-layers", type=int, required=True)
parser.add_argument("--hidden-width", type=int, required=True)
parser.add_argument("--dropout", type=float, required=True)
parser.add_argument("--optimizer", choices=["Adam", "RMSprop", "SGD"], required=True)
parser.add_argument("--learning-rate", type=float, required=True)
parser.add_argument("--epochs", type=int, default=5)
parser.add_argument("--batch-size", type=int, default=128)
parser.add_argument("--seed", type=int, default=23)
args = parser.parse_args()

torch.manual_seed(args.seed)
tracking.init()

try:
    device = torch.device("cuda")
    dataset = MNIST(
        root="/tmp/mnist", train=True, download=True, transform=ToTensor()
    )
    train_set, validation_set = random_split(
        dataset, [50_000, 10_000],
        generator=torch.Generator().manual_seed(42),
    )
    train_loader = DataLoader(
        train_set,
        batch_size=args.batch_size,
        shuffle=True,
        generator=torch.Generator().manual_seed(args.seed),
        pin_memory=True,
    )
    validation_loader = DataLoader(
        validation_set, batch_size=512, pin_memory=True
    )

    layers = [nn.Flatten()]
    in_features = 28 * 28
    for _ in range(args.n_layers):
        layers.extend([
            nn.Linear(in_features, args.hidden_width),
            nn.ReLU(),
            nn.Dropout(args.dropout),
        ])
        in_features = args.hidden_width
    layers.append(nn.Linear(in_features, 10))
    model = nn.Sequential(*layers).to(device)

    optimizer_class = getattr(torch.optim, args.optimizer)
    optimizer = optimizer_class(model.parameters(), lr=args.learning_rate)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(1, args.epochs + 1):
        model.train()
        for images, labels in train_loader:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()

        model.eval()
        correct = 0
        with torch.inference_mode():
            for images, labels in validation_loader:
                predictions = model(
                    images.to(device, non_blocking=True)
                ).argmax(dim=1)
                correct += (
                    predictions == labels.to(device, non_blocking=True)
                ).sum().item()

        validation_accuracy = correct / len(validation_set)
        tracking.log_metrics(
            step=epoch, validation_accuracy=validation_accuracy
        )

    checkpoint = Path(tracking.get_outputs_path("model.pt"))
    checkpoint.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {"state_dict": model.state_dict(), "parameters": vars(args)},
        checkpoint,
    )
    tracking.log_file_ref(path=str(checkpoint), name="model", is_input=False)
    tracking.log_outputs(
        async_req=False, validation_accuracy=validation_accuracy
    )
finally:
    tracking.end()

The MNIST training dataset supplies the 50,000 training and 10,000 validation examples. The split stays fixed across trials. The separate test dataset is left out of the search.

Each trial constructs a fresh model and optimizer. Its final layer returns logits, which CrossEntropyLoss accepts directly.

The two tracking calls serve different purposes. log_metrics records the validation curve for comparison; log_outputs returns the final epoch's accuracy under the objective name declared in the matrix. That output is what the tuner uses. The checkpoint and its architecture parameters are saved with the trial's artifacts.

Run the study and compare trials

From the folder containing all three files, submit the study to your project:

polyaxon run --project=acme/hpo -f study.yaml -u

Replace acme/hpo with your project. The upload option sends the current directory to the run's uploads folder, matching the component's working directory.

Open the study's child runs and use run comparison to compare their parameters and validation curves. Sort successful training trials by validation_accuracy in descending order, then inspect the leading configurations and their model artifacts.

Illustrative validation-accuracy curves for three trials over five epochs, with no numeric accuracy scale. Trial A peaks earlier, while Trial B has the highest final-epoch accuracy used to rank the trials.

Illustrative curves, not results from this MNIST study. Compare the full validation history, but rank trials by the final-epoch output: an earlier peak does not change the objective passed to TPE.

Keep the five-epoch training budget and validation split fixed while comparing this study. The leading configuration is the strongest within this budget; a longer training run may change the ranking. Before choosing a model, confirm a shortlist across training seeds, then evaluate the selected configuration on the held-out test set.

What about pruning?

The training script reports validation accuracy to Polyaxon after each epoch. Reporting a metric does not stop a trial on its own: this example trains each configuration for five epochs unless it fails or reaches its timeout.

Polyaxon also supports early-stopping rules configured on the matrix. A metric condition without a policy can end the whole study; median, truncation, and difference policies can target underperforming runs. Configure these separately when needed. This example leaves them disabled so every trial receives the same five-epoch budget.