Skip to content

CIFAR-10

Introduction

Here we will walk through an example project for training a model to classify images from the CIFAR10 dataset. The full source lives in the repo under examples/cifar10/. The flow is end-to-end py-gen-ml. You author a protobuf schema, generate typed config tooling, then load YAML, apply CLI overrides, and optionally sweep. The training loop stays ordinary Python. This combines most of the concepts from the guides so far. It offers a simple model, data and optimizer, and shows how to:

  • Use a variety of special attributes to define the config schema
  • Load a config from a yaml file
  • Apply CLI arguments to the config
  • Load a sweep file to sample from a set of configurations

Schema

The schema is defined in a proto file. This is the file that we will use to train the model. It looks like this:

The model

The model definition is a simple convolutional neural network followed by a multi-layer perceptron. In proto definition, it looks like this:

// Model configuration
message Model {
    // Conv blocks
    ConvNet conv_net = 1;
    // MLP head
    MLP head = 2;
}

// Convolutional neural network configuration
message ConvNet {
    // Conv layer configuration
    ConvBlock block = 1;
    // Number of layers
    uint32 num_layers = 2 [(pgml.default).uint32 = 2];
}

// Multi-layer perceptron configuration
message MLP {
    // Linear layer configuration
    LinearBlock block = 1;
    // Number of layers
    uint32 num_layers = 2 [(pgml.default).uint32 = 2];
}

// Convolutional layer configuration
message ConvBlock {
    option (pgml.factory) = "cifar10.modules.ConvBlock";
    // Number of output channels
    uint32 out_channels = 1  [(pgml.default).uint32 = 128];
    // Square kernel size
    uint32 kernel_size = 2 [(pgml.default).uint32 = 3];
    // Square pool size
    uint32 pool_size = 3 [(pgml.default).uint32 = 2];
    // Activation function
    Activation activation = 4 [(pgml.default).enum = "GELU"];
}

// Linear layer configuration
message LinearBlock {
    option (pgml.factory) = "cifar10.modules.LinearBlock";
    // Number of output features
    uint32 out_features = 1 [(pgml.default).uint32 = 128];
    // Activation function
    Activation activation = 2 [(pgml.default).enum = "GELU"];
}

// Activation function
enum Activation {
    // GELU activation
    GELU = 0;
    // ReLU activation
    RELU = 1;
}

Concepts that we've used here are:

  • Nesting: we define a Model message that contains two nested messages: ConvNet and MLP which in turn contain blocks of linear and convolutional layers.
  • pgml.factory: This is a special attribute that tells py-gen-ml that with these values we can build an instance of the given class.
  • pgml.default: This is a special attribute that tells py-gen-ml to use a default value if the field is not set.
  • Enums: we define an enum for the activation function.

The data

We define a Data message that contains the batch size and the number of epochs to train. In proto definition, it looks like this:

// Data configuration
message Data {
    // Batch size for a single GPU
    uint32 batch_size = 1 [(pgml.default).uint32 = 32];
    // Number of epochs to train
    uint32 num_epochs = 2 [(pgml.default).uint32 = 10];
}

The optimizer

We define an Optimizer message that contains the learning rate and the decay rate. In proto definition, it looks like this:

// Optimizer configuration
message Optimizer {
    // Learning rate
    float learning_rate = 1 [(pgml.default).float = 1e-4];
    // Decay rate
    float beta1 = 2 [(pgml.default).float = 0.99];
}

The project

We define a Project message that contains the model, the optimizer, and the data. In proto definition, it looks like this:

syntax = "proto3";

package cifar10;

import "py_gen_ml/extensions.proto";


// Global configuration
message Project {
    option (pgml.cli) = {
        enable: true
        arg: {name: "conv_activation", path: "net.conv_net.block.activation"}
        arg: {name: "head_activation", path: "net.head.block.activation"}
        arg: {name: "num_conv_layers", path: "net.conv_net.num_layers"}
        arg: {name: "num_mlp_layers", path: "net.head.num_layers"}
    };
    // Model configuration
    Model net = 1;
    // Optimizer configuration
    Optimizer optimizer = 2;
    // Data configuration
    Data data = 3;
}
As you can see, there are a couple of arg references to ensure we can propagate values via the command line for fields that have the same name and that are deeply nested.

The entrypoint

To launch the training, we create a function that can do a range of things:

  1. Load a yaml config file
  2. Load multiple yaml config files and merge them
  3. Load a yaml sweep file to generate patches and apply them to the config using Optuna
  4. Apply CLI arguments to the config
  5. Whatever the input is, train the model and return the accuracy

Even though this sounds like a lot, with pg-gen-ml it is actually quite easy to do. The main function now becomes:

Entrypoint
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
@pgml.pgml_cmd(app=app)
def main(
    config_paths: List[str] = typer.Option(..., help='Paths to config files'),
    sweep_paths: List[str] = typer.Option(default_factory=list, help='Paths to sweep files'),
    cli_args: cli_args.ProjectArgs = typer.Option(...),
) -> None:
    config = base.Project.from_yaml_files(config_paths)
    config = config.apply_cli_args(cli_args)

    if len(sweep_paths) == 0:
        train_model(config, trial=None)
        return

    sweep_config = sweep.ProjectSweep.from_yaml_files(sweep_paths)

    def objective(trial: optuna.Trial) -> float:
        sampler = pgml.OptunaSampler(trial=trial)
        patch = sampler.sample(sweep_config)
        accuracy = train_model(project=config.merge(patch), trial=trial)
        return accuracy

    study = optuna.create_study(direction='maximize')
    study.optimize(objective, n_trials=100)

Let's break this down a bit more.

  • At line 7, we load the project config from one or more yaml files. This is where files will be merged before they are passed to the Pydantic model validator.
  • At line 8, we apply the CLI arguments to the config.
  • At line 10, we check if there are any sweep files. If there are, we load them at line 14. If there are no sweep files, we simply train the model and return the accuracy.
  • At line 16-20, we define an objective function for Optuna to use. This is the function that will be called to train the model and return the accuracy. Note that we are using the exact same train_model function that we used in line 10 earlier.
  • At line 22 and 23, we create an Optuna study and optimize the objective function.

The configuration

The configuration is defined in a yaml file. This is the file that we will use to train the model. It looks like this:

Config
# yaml-language-server: $schema=schemas/project.json
net:
  conv_net:
    block:
      activation: GELU
      kernel_size: 3
      out_channels: 32
      pool_size: 2
  head:
    block:
      activation: GELU
      out_features: 128
    num_layers: 3
optimizer:
  beta1: 0.99
  learning_rate: 0.0001
data:
  batch_size: 32
  num_epochs: 10

Sweep configuration

We define the following sweep configuration:

Sweep
# yaml-language-server: $schema=schemas/project.json
optimizer:
  beta1:
    low: 0.9
    high: 0.99
  learning_rate:
    log_low: 1e-5
    log_high: 1e-3

This is a sweep over the learning rate and beta1 parameters. We'll use it below when launching Optuna trials.

Example code overview

Hand-written training code lives under examples/cifar10/src/cifar10/. Generated config tooling lives under examples/cifar10/src/pgml_out/. Rough map:

Piece Path Role
Schema config.proto Source of truth for the project config
Generated models pgml_out/config_{base,patch,sweep,cli_args}.py Load YAML, merge patches, sample sweeps, parse CLI
Modules modules.py Conv / MLP building blocks and Model.from_config
Data data.py CIFAR-10 transforms and DataLoaders
Trainer train.py (Trainer) Train / eval loop, optional Optuna trial reporting
Train function train.py (train_model) Wire config → model, data, optimizer, trainer
Entrypoint train.py (main) @pgml.pgml_cmd CLI: configs, sweeps, CLI overrides

Modules

Modules
from typing import Self

import pgml_out.config_base as base
import torch

CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']


def create_activation(activation: base.Activation) -> torch.nn.Module:
    if activation == base.Activation.GELU:
        return torch.nn.GELU()
    elif activation == base.Activation.RELU:
        return torch.nn.ReLU()
    else:
        raise ValueError(f'Invalid activation function: {activation}')


class LinearBlock(torch.nn.Module):
    """Linear layer with activation configuration"""

    def __init__(self, out_features: int, activation: base.Activation) -> None:
        super().__init__()
        self.linear = torch.nn.LazyLinear(out_features=out_features)
        self.activation = create_activation(activation)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.activation(self.linear(x))


class ConvBlock(torch.nn.Module):
    """Convolutional layer with activation configuration"""

    def __init__(self, out_channels: int, kernel_size: int, pool_size: int, activation: base.Activation) -> None:
        super().__init__()
        self.conv = torch.nn.LazyConv2d(out_channels=out_channels, kernel_size=kernel_size)
        self.pool = torch.nn.MaxPool2d(kernel_size=pool_size)
        self.activation = create_activation(activation)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.activation(self.pool(self.conv(x)))


class ConvNet(torch.nn.Module):

    def __init__(self, block: base.ConvBlock, num_layers: int) -> None:
        super().__init__()
        self.layers = [block.build() for _ in range(num_layers)]
        self.net = torch.nn.Sequential(*self.layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MLP(torch.nn.Module):

    def __init__(self, block: base.LinearBlock, num_layers: int) -> None:
        super().__init__()
        self.layers = [block.build() for _ in range(num_layers)]
        self.net = torch.nn.Sequential(*self.layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class Model(torch.nn.Module):

    def __init__(self, conv_net: ConvNet, head: MLP) -> None:
        super().__init__()
        self.conv_net = conv_net
        self.head = head
        self.class_logits = torch.nn.LazyLinear(out_features=len(CLASSES),)
        self.flatten = torch.nn.Flatten()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv_net(x)
        x = self.head(x)
        x = self.flatten(x)
        x = self.class_logits(x)
        return x

    @classmethod
    def from_config(cls, config: base.Model) -> Self:
        conv_net = ConvNet(
            block=config.conv_net.block,
            num_layers=config.conv_net.num_layers,
        )
        head = MLP(
            block=config.head.block,
            num_layers=config.head.num_layers,
        )
        return cls(conv_net, head)

Data

Data
import os
from typing import Any

import torch
import torchvision
import torchvision.transforms as transforms


def get_transform() -> transforms.Compose:
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ],)


def get_data_loader(transform: transforms.Compose, batch_size: int, train: bool) -> torch.utils.data.DataLoader[Any]:
    trainset = torchvision.datasets.CIFAR10(
        root=f"{os.environ['HOME']}/data/torchvision",
        train=train,
        download=train,
        transform=transform,
    )
    return torch.utils.data.DataLoader[Any](
        trainset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        prefetch_factor=4,
    )

Trainer

Trainer
class Trainer:

    def __init__(
        self,
        model: torch.nn.Module,
        train_loader: torch.utils.data.DataLoader[Any],
        test_loader: torch.utils.data.DataLoader[Any],
        criterion: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        accuracy_metric_train: torchmetrics.classification.MulticlassAccuracy,
        accuracy_metric_test: torchmetrics.classification.MulticlassAccuracy,
        train_loss_metric: torchmetrics.MeanMetric,
        test_loss_metric: torchmetrics.MeanMetric,
        trial: Optional[optuna.Trial] = None,
    ) -> None:
        self._train_loader = train_loader
        self._test_loader = test_loader
        self._model = model
        self._criterion = criterion
        self._optimizer = optimizer
        self._accuracy_metric_train = accuracy_metric_train
        self._accuracy_metric_test = accuracy_metric_test
        self._train_loss_metric = train_loss_metric
        self._test_loss_metric = test_loss_metric
        self._trial = trial

    def train(self, num_epochs: int) -> float:
        device = get_device()
        step = 0
        for epoch in tqdm.trange(num_epochs, position=0, desc='Epoch'):
            batch_bar = tqdm.tqdm(self._train_loader, position=1, desc='Batch')
            for inputs, labels in batch_bar:
                inputs = inputs.to(device)
                labels = labels.to(device)

                self._optimizer.zero_grad()
                outputs = self._model(inputs)
                loss = self._criterion(outputs, labels)
                loss.backward()
                self._optimizer.step()

                self._train_loss_metric.update(loss)
                self._accuracy_metric_train.update(outputs, labels)
                step += 1

                if step % 10 == 0:
                    batch_bar.set_postfix(
                        loss=self._train_loss_metric.compute().item(),
                        accuracy=self._accuracy_metric_train.compute().item(),
                    )
                    self._train_loss_metric.reset()
                    self._accuracy_metric_train.reset()

            self._evaluate()
            if self._trial is not None:
                self._trial.report(self._accuracy_metric_test.compute().item(), epoch)
        return self._accuracy_metric_test.compute().item()

    @torch.inference_mode()
    def _evaluate(self) -> None:
        self._model.eval()
        self._accuracy_metric_test.reset()
        self._test_loss_metric.reset()
        for images, labels in tqdm.tqdm(self._test_loader, position=1, desc='Evaluating'):
            images = images.to(get_device())
            labels = labels.to(get_device())
            outputs = self._model(images)
            loss = self._criterion(outputs, labels)
            self._test_loss_metric.update(loss)
            self._accuracy_metric_test.update(outputs, labels)
        print(f'Test accuracy: {self._accuracy_metric_test.compute().item()}')
        print(f'Test loss: {self._test_loss_metric.compute().item()}')
        self._model.train()

Train function

The train function that instantiates all the components and calls the trainer:

Train function
def train_model(project: base.Project, trial: typing.Optional[optuna.Trial] = None) -> float:
    rich.print(project)

    transform = get_transform()
    train_loader = get_data_loader(transform=transform, batch_size=project.data.batch_size, train=True)
    test_loader = get_data_loader(transform=transform, batch_size=project.data.batch_size, train=False)

    device = get_device()
    print(f'device {device}')

    model = Model.from_config(config=project.net).to(device)
    path = pathlib.Path(f"{os.environ['HOME']}/gen_ml/logs/{uuid.uuid4()}")
    print(f'Storing logs at {path}')

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(params=model.parameters(), lr=project.optimizer.learning_rate)
    accuracy_metric_train = get_accuracy_metric(num_classes=len(CLASSES))
    accuracy_metric_test = get_accuracy_metric(num_classes=len(CLASSES))
    train_loss_metric = torchmetrics.MeanMetric().to(device)
    test_loss_metric = torchmetrics.MeanMetric().to(device)
    trainer = Trainer(
        model=model,
        train_loader=train_loader,
        test_loader=test_loader,
        criterion=criterion,
        optimizer=optimizer,
        accuracy_metric_train=accuracy_metric_train,
        accuracy_metric_test=accuracy_metric_test,
        train_loss_metric=train_loss_metric,
        test_loss_metric=test_loss_metric,
        trial=trial,
    )
    accuracy = trainer.train(num_epochs=project.data.num_epochs)
    return accuracy

How to run

All commands below assume you are in the example project directory:

cd examples/cifar10

Install the example (pulls in py-gen-ml from the workspace plus Torch):

uv sync

Inspect the CLI

uv run python src/cifar10/train.py --help

This prints the generated options, including nested field overrides:

 Usage: train.py [OPTIONS]

╭─ Options ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ *  --config-paths              TEXT         Path to config file [default: None] [required]                                             │
│    --sweep-paths               TEXT         Type of config to use [default: <class 'list'>]                                            │
│    --out-channels              INTEGER      Number of output channels. Maps to 'net.conv_net.block.out_channels' [default: None]       │
│    --kernel-size               INTEGER      Square kernel size. Maps to 'net.conv_net.block.kernel_size' [default: None]               │
│    --pool-size                 INTEGER      Square pool size. Maps to 'net.conv_net.block.pool_size' [default: None]                   │
│    --out-features              INTEGER      Number of output features. Maps to 'net.head.block.out_features' [default: None]           │
│    --batch-size                INTEGER      Batch size for a single GPU. Maps to 'data.batch_size' [default: None]                     │
│    --num-epochs                INTEGER      Number of epochs to train. Maps to 'data.num_epochs' [default: None]                       │
│    --learning-rate             FLOAT        Learning rate. Maps to 'optimizer.learning_rate' [default: None]                           │
│    --beta1                     FLOAT        Decay rate. Maps to 'optimizer.beta1' [default: None]                                      │
│    --conv-activation           [gelu|relu]  Activation function. Maps to 'net.conv_net.block.activation' [default: None]               │
│    --head-activation           [gelu|relu]  Activation function. Maps to 'net.head.block.activation' [default: None]                   │
│    --num-conv-layers           INTEGER      Number of layers. Maps to 'net.conv_net.num_layers' [default: None]                        │
│    --num-mlp-layers            INTEGER      Number of layers. Maps to 'net.head.num_layers' [default: None]                            │
│    --install-completion                     Install completion for the current shell.                                                  │
│    --show-completion                        Show completion for the current shell, to copy it or customize the installation.           │
│    --help                                   Show this message and exit.                                                                │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

Simple run (base config only)

Train once with the default YAML:

uv run python src/cifar10/train.py \
    --config-paths configs/base/default.yaml

CLI overrides

Same base config, but override a few fields from the command line (no second YAML needed):

uv run python src/cifar10/train.py \
    --config-paths configs/base/default.yaml \
    --learning-rate 1e-3 \
    --batch-size 64 \
    --num-epochs 2

Hyperparameter sweep

Run an Optuna study over the learning rate and beta1 search space. Each trial merges a sampled patch onto the base config and calls the same train_model:

uv run python src/cifar10/train.py \
    --config-paths configs/base/default.yaml \
    --sweep-paths configs/sweep/lr_beta1.yaml

Sweep plus CLI overrides

CLI args still apply on top of the base config before sampling. Useful for fixing something like epoch count while sweeping the optimizer:

uv run python src/cifar10/train.py \
    --config-paths configs/base/default.yaml \
    --sweep-paths configs/sweep/lr_beta1.yaml \
    --num-epochs 1 \
    --batch-size 128

Tip

Pass multiple --config-paths to merge base + patch YAML files before training or sweeping. Later files overlay earlier ones the same way as elsewhere in py-gen-ml.

Conclusion

This example is a small but complete training project built around a protobuf schema. You write config.proto and the YAML. py-gen-ml generates the typed base / patch / sweep / CLI layer. train.py stays ordinary PyTorch. The same entrypoint covers a single run, CLI experiments, and Optuna sweeps without forking the train loop.