helios-ml 0.1.0__tar.gz

helios-ml-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,46 @@
+name: "Publish"
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment:
+      name: testpypi
+      url: https://pypi.org/project/helios-ml
+    permissions:
+      id-token: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        id: setup_python
+        with:
+          python-version: "3.11"
+
+      - name: Cache virtualenv
+        uses: actions/cache@v3
+        with:
+          key: venv-${{ runner.os }}-${{ steps.setup_python.outputs.python-version}}-${{ hashFiles('requirements/ci.txt') }}
+          path: venv
+
+      - name: Install dependencies
+        run: |
+          python -m venv venv
+          source venv/bin/activate
+          python -m pip install -r requirements/ci.txt
+          echo "$VIRTUAL_ENV/bin" >> $GITHUB_PATH
+          echo "VIRTUAL_ENV=$VIRTUAL_ENV" >> $GITHUB_ENV
+
+      - name: Build Helios
+        run: |
+          python -m build
+      - name: Publish Helios
+        uses: pypa/gh-action-pypi-publish@release/v1

helios-ml-0.1.0/.github/workflows/tests.yml

@@ -0,0 +1,66 @@
+name: "Tests"
+
+on:
+  push:
+    branches: [ master ]
+    paths-ignore:
+      - 'README.md'
+      - '.pre-commit-config.yaml'
+      - 'data/**'
+      - "LICENSE"
+  pull_request:
+    branches: [ master ]
+    paths-ignore:
+      - 'README.md'
+      - 'LICENSE'
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        id: setup_python
+        with:
+          python-version: "3.11"
+
+      - name: Cache virtualenv
+        uses: actions/cache@v3
+        with:
+          key: venv-${{ runner.os }}-${{ steps.setup_python.outputs.python-version}}-${{ hashFiles('requirements/ci.txt') }}
+          path: venv
+
+      - name: Install dependencies
+        run: |
+          python -m venv venv
+          source venv/bin/activate
+          python -m pip install -r requirements/ci.txt
+          echo "$VIRTUAL_ENV/bin" >> $GITHUB_PATH
+          echo "VIRTUAL_ENV=$VIRTUAL_ENV" >> $GITHUB_ENV
+
+      - name: Install Helios
+        shell: bash
+        run: |
+          pip install .
+
+      - name: Ruff
+        run: |
+          ruff check src/helios
+          ruff check test
+          ruff check examples/cifar10
+
+      - name: Mypy
+        run: |
+          mypy src/helios
+          mypy test
+          mypy examples/cifar10
+
+      - name: Pytest
+        run: |
+          python -m pytest

helios-ml-0.1.0/.gitignore

@@ -0,0 +1,20 @@
+# Virtual envs.
+venv/
+.venv/
+
+# Distribution/packaging
+dist/
+build/
+*.egg-info/
+
+__pycache__/
+.mypy_cache/
+.pytest_cache/
+.ruff_cache/
+
+# Generated folders from examples
+*.pth
+examples/**/runs
+examples/**/chkpt
+examples/**/logs
+examples/**/data

helios-ml-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: c4a0b883114b00d8d76b479c820ce7950211c99b # frozen: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+      - id: debug-statements
+      - id: check-ast
+      - id: mixed-line-ending
+        args: [--fix=lf]
+      - id: check-yaml
+        args: [--allow-multiple-documents]
+      - id: check-json
+      - id: check-added-large-files
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.3.4
+    hooks:
+      - id: ruff-format
+
+  - repo: local
+    hooks:
+      - id: generate_requirements.py
+        name: generate_requirements.py
+        language: system
+        entry: python tools/generate_requirements.py
+        files: "pyproject.toml|requirements/.*\\.txt|tools/generate_requirements.py"

helios-ml-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, Mauricio A Rovira Galvez
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

helios-ml-0.1.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.1
+Name: helios-ml
+Version: 0.1.0
+Summary: A Torch-based package for training AI networks
+Author: Mauricio A. Rovira Galvez
+Project-URL: Homepage, https://github.com/marovira/pyro-ml
+Project-URL: Issues, https://github.com/marovira/pyro-ml/issues
+Requires-Python: >=3.11
+Description-Content-Type: text/x-rst
+License-File: LICENSE
+Requires-Dist: tqdm>=4.66.2
+Requires-Dist: opencv-python>=4.9.0.80
+Requires-Dist: protobuf!=4.24.0,<5.0.0,>=3.19.6
+Requires-Dist: tensorboard>=2.16.2
+Requires-Dist: torch>=2.2.1
+Requires-Dist: torchvision>=0.17.1
+Requires-Dist: onnx>=1.16.0
+Requires-Dist: onnxruntime>=1.17.1
+Requires-Dist: matplotlib>=3.8.4
+Provides-Extra: dev
+Requires-Dist: mypy>=1.8.0; extra == "dev"
+Requires-Dist: ruff>=0.3.4; extra == "dev"
+Requires-Dist: pytest>=8.1.1; extra == "dev"
+Requires-Dist: pre-commit>=3.6.2; extra == "dev"
+Requires-Dist: types-Pillow>=10.2.0.20240311; extra == "dev"
+Requires-Dist: types-tqdm>=4.66.0.20240106; extra == "dev"
+Requires-Dist: build>=1.2.1; extra == "dev"
+Requires-Dist: twine>=5.0.0; extra == "dev"
+Provides-Extra: ci
+Requires-Dist: mypy>=1.8.0; extra == "ci"
+Requires-Dist: ruff>=0.3.4; extra == "ci"
+Requires-Dist: pytest>=8.1.1; extra == "ci"
+Requires-Dist: types-Pillow>=10.2.0.20240311; extra == "ci"
+Requires-Dist: types-tqdm>=4.66.0.20240106; extra == "ci"
+Requires-Dist: build>=1.2.1; extra == "ci"
+
+<a id="top"></a>
+![HELIOS logo](data/logo/logo-transparent.png)
+
+[![Generic badge](https://img.shields.io/badge/License-BSD3-blue)](LICENSE)
+[![Static Badge](https://img.shields.io/badge/Python-3.11%2B-red?logoColor=red)](https://www.python.org/downloads/release/python-3110/)
+[![Tests](https://github.com/marovira/helios-ml/actions/workflows/tests.yml/badge.svg)](https://github.com/marovira/helios-ml/actions/workflows/tests.yml)
+
+## What is Helios?
+
+Named after Greek god of the sun, Helios is a light-weight package for training ML
+networks built on top of PyTorch. It is designed to abstract all of the "boiler-plate"
+code involved with training. Specifically, it wraps the following common patterns:
+
+* Creation of the dataloaders.
+* Initialization of CUDA, PyTorch, and random number states.
+* Initialization for distributed training.
+* Training, validation, and testing loops.
+* Saving and loading checkpoints.
+* Exporting to ONNX.
+
+It is important to note that Helios is **not** a fully fledged training environment similar
+to [Pytorch Lightning](https://github.com/Lightning-AI/pytorch-lightning). Instead, Helios
+is focused on providing a simple and straight-forward interface that abstracts most of the
+common code patterns while retaining the ability to be easily overridden to suit the
+individual needs of each training scheme.
+
+## Main Features
+
+Helios offers the following functionality out of the box:
+
+1. Resume training: Helios has been built with the ability to resume training if it is
+   paused. Specifically, Helios will ensure that the behaviour of the trained model is
+   *identical* to the one it would've had if it had been trained without pauses.
+2. Automatic detection of multi-GPU environments for distributed training. In addition,
+   Helios also supports training using `torchrun` and will automatically handle the
+   initialisation and clean up of the distributed state. It will also correctly set the
+   devices and maps to ensure weights are mapped tot he correct location.
+3. Registries for creation of arbitrary types. These include: networks, loss functions,
+   optimizers, schedulers, etc.
+4. Correct handling of logging when doing distributed training (even over multiple nodes).
+
+## Installation
+
+You can install Helios using Pip as follows:
+
+```sh
+pip install helios-ml
+```
+
+## Documentation
+
+Documentation coming soon!
+
+## Contributing
+
+There are three ways in which you can contribute to Helios:
+
+* If you find a bug, please open an issue. Similarly, if you have a question
+  about how to use it, or if something is unclear, please post an issue so it
+  can be addressed.
+* If you have a fix for a bug, or a code enhancement, please open a pull
+  request. Before you submit it though, make sure to abide by the rules written
+  below.
+* If you have a feature proposal, you can either open an issue or create a pull
+  request. If you are submitting a pull request, it must abide by the rules
+  written below. Note that any new features need to be approved by me.
+
+If you are submitting a pull request, the guidelines are the following:
+
+1. Ensure that your code follows the standards and formatting of Helios. The coding
+   standards and formatting are enforced through the Ruff Linter and Formatter. Any
+   changes that do not abide by these rules will be rejected. It is your responsibility to
+   ensure that both Ruff and Mypy linters pass.
+2. Ensure that *all* unit tests are working prior to submitting the pull
+   request. If you are adding a new feature that has been approved, it is your
+   responsibility to provide the corresponding unit tests (if applicable).
+
+## License
+
+Helios is published under the BSD-3 license and can be viewed [here](LICENSE).

helios-ml-0.1.0/README.md

@@ -0,0 +1,80 @@
+<a id="top"></a>
+![HELIOS logo](data/logo/logo-transparent.png)
+
+[![Generic badge](https://img.shields.io/badge/License-BSD3-blue)](LICENSE)
+[![Static Badge](https://img.shields.io/badge/Python-3.11%2B-red?logoColor=red)](https://www.python.org/downloads/release/python-3110/)
+[![Tests](https://github.com/marovira/helios-ml/actions/workflows/tests.yml/badge.svg)](https://github.com/marovira/helios-ml/actions/workflows/tests.yml)
+
+## What is Helios?
+
+Named after Greek god of the sun, Helios is a light-weight package for training ML
+networks built on top of PyTorch. It is designed to abstract all of the "boiler-plate"
+code involved with training. Specifically, it wraps the following common patterns:
+
+* Creation of the dataloaders.
+* Initialization of CUDA, PyTorch, and random number states.
+* Initialization for distributed training.
+* Training, validation, and testing loops.
+* Saving and loading checkpoints.
+* Exporting to ONNX.
+
+It is important to note that Helios is **not** a fully fledged training environment similar
+to [Pytorch Lightning](https://github.com/Lightning-AI/pytorch-lightning). Instead, Helios
+is focused on providing a simple and straight-forward interface that abstracts most of the
+common code patterns while retaining the ability to be easily overridden to suit the
+individual needs of each training scheme.
+
+## Main Features
+
+Helios offers the following functionality out of the box:
+
+1. Resume training: Helios has been built with the ability to resume training if it is
+   paused. Specifically, Helios will ensure that the behaviour of the trained model is
+   *identical* to the one it would've had if it had been trained without pauses.
+2. Automatic detection of multi-GPU environments for distributed training. In addition,
+   Helios also supports training using `torchrun` and will automatically handle the
+   initialisation and clean up of the distributed state. It will also correctly set the
+   devices and maps to ensure weights are mapped tot he correct location.
+3. Registries for creation of arbitrary types. These include: networks, loss functions,
+   optimizers, schedulers, etc.
+4. Correct handling of logging when doing distributed training (even over multiple nodes).
+
+## Installation
+
+You can install Helios using Pip as follows:
+
+```sh
+pip install helios-ml
+```
+
+## Documentation
+
+Documentation coming soon!
+
+## Contributing
+
+There are three ways in which you can contribute to Helios:
+
+* If you find a bug, please open an issue. Similarly, if you have a question
+  about how to use it, or if something is unclear, please post an issue so it
+  can be addressed.
+* If you have a fix for a bug, or a code enhancement, please open a pull
+  request. Before you submit it though, make sure to abide by the rules written
+  below.
+* If you have a feature proposal, you can either open an issue or create a pull
+  request. If you are submitting a pull request, it must abide by the rules
+  written below. Note that any new features need to be approved by me.
+
+If you are submitting a pull request, the guidelines are the following:
+
+1. Ensure that your code follows the standards and formatting of Helios. The coding
+   standards and formatting are enforced through the Ruff Linter and Formatter. Any
+   changes that do not abide by these rules will be rejected. It is your responsibility to
+   ensure that both Ruff and Mypy linters pass.
+2. Ensure that *all* unit tests are working prior to submitting the pull
+   request. If you are adding a new feature that has been approved, it is your
+   responsibility to provide the corresponding unit tests (if applicable).
+
+## License
+
+Helios is published under the BSD-3 license and can be viewed [here](LICENSE).

helios-ml-0.1.0/examples/cifar10/cifar10.py

@@ -0,0 +1,297 @@
+import os
+import pathlib
+import typing
+
+import torch
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms.v2 as T
+from torch import nn
+
+import helios.core as hlc
+import helios.data as hld
+import helios.model as hlm
+import helios.optim as hlo
+import helios.trainer as hlt
+from helios.core import logging
+
+
+class CIFARDataModule(hld.DataModule):
+    """
+    Example datamodule class built with CIFAR10.
+
+    Here you can see how to setup a datamodule that requires data to be downloaded using
+    torchvision's CIFAR10 dataset. The code is adapted from PyTorch's "Training a
+    Classifier" tutorial:
+    https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
+
+    Args:
+        root (pathlib.Path): the root where the data will be downloaded to.
+    """
+
+    def __init__(self, root: pathlib.Path) -> None:
+        """Construct the datamodule."""
+        super().__init__()
+        self._root = root / "data"
+
+    def prepare_data(self) -> None:
+        """Download all necessary data."""
+        torchvision.datasets.CIFAR10(root=self._root, train=True, download=True)
+        torchvision.datasets.CIFAR10(root=self._root, train=False, download=True)
+
+    def setup(self) -> None:
+        """Create the datasets."""
+        # Use the ToTensor transform from Pyro to automate the conversion from images to
+        # tensors.
+        transforms = T.Compose(
+            [hld.transforms.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+        )
+        params = hld.DataLoaderParams()
+        params.batch_size = 4
+        params.shuffle = True
+        params.num_workers = 2
+        params.drop_last = True
+        self._train_dataset = self._create_dataset(
+            torchvision.datasets.CIFAR10(
+                root=self._root, train=True, download=False, transform=transforms
+            ),
+            params,
+        )
+
+        # The dataloader params are copied when the dataset is created, so we can safely
+        # change the options for the validation dataset without interfering with the ones
+        # for training.
+        params.drop_last = False
+        params.shuffle = False
+        self._valid_dataset = self._create_dataset(
+            torchvision.datasets.CIFAR10(
+                root=self._root, train=False, download=False, transform=transforms
+            ),
+            params,
+        )
+
+
+class Net(nn.Module):
+    """
+    Example image classifier.
+
+    The code is taken from PyTorch's Training a Classifier tutorial:
+    https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
+    """
+
+    def __init__(self):
+        """Create the classifier."""
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Compute the label for the given image."""
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = torch.flatten(x, 1)  # flatten all dimensions except batch
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
+class ClassifierModel(hlm.Model):
+    """
+    Example model class for training the classifier.
+
+    Here you can see how to setup the model class and some of the basic functionality that
+    is available. The code is adapted from PyTorch's "Training a Classifier" tutorial:
+    https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
+    """
+
+    def __init__(self) -> None:
+        """Create the model."""
+        super().__init__("classifier")
+
+    def setup(self, fast_init: bool = False) -> None:
+        """Create the network, loss, and optimizer."""
+        # Note that when we create the network and loss function, we immediately move them
+        # to the current device, which has been set by the trainer.
+        self._net = Net().to(self.device)
+        self._criterion = nn.CrossEntropyLoss().to(self.device)
+
+        # Note that SGD is shipped as part of the default optimizers from Pyro, so we can
+        # directly request it from create_optimizer instead of building it ourselves.
+        self._optimizer = hlo.create_optimizer(
+            "SGD", self._net.parameters(), lr=0.001, momentum=0.9
+        )
+
+    def load_state_dict(
+        self, state_dict: dict[str, typing.Any], fast_init: bool = False
+    ) -> None:
+        """Restore the model from a saved checkpoint."""
+        # Note that we don't have to re-map the weights ourselves. They have already been
+        # re-mapped for us by the trainer when it loaded the checkpoint.
+        self._net.load_state_dict(state_dict["net"])
+        self._criterion.load_state_dict(state_dict["criterion"])
+        self._optimizer.load_state_dict(state_dict["optimizer"])
+
+    def state_dict(self) -> dict[str, typing.Any]:
+        """Return the state dict of the model."""
+        return {
+            "net": self._net.state_dict(),
+            "criterion": self._criterion.state_dict(),
+            "optimizer": self._optimizer.state_dict(),
+        }
+
+    def train(self) -> None:
+        """Set the model to train."""
+        # If we had more networks, we would shift them to training mode here.
+        self._net.train()
+
+    def on_training_start(self) -> None:
+        """Perform steps before training starts."""
+        tb_logger = hlc.get_from_optional(logging.get_tensorboard_writer())
+
+        x = torch.randn((1, 3, 32, 32)).to(self.device)
+        tb_logger.add_graph(self._net, x)
+
+    def train_step(self, batch: typing.Any, state: hlt.TrainingState) -> None:
+        """Forward and backward training passes."""
+        # Due to the simplicity of the code, we do both the forward and backward passes in
+        # the training step, but you could also split it between the train_step and
+        # on_training_batch_end if your setup is more complex.
+        inputs, labels = batch
+        inputs = inputs.to(self.device)
+        labels = labels.to(self.device)
+
+        self._optimizer.zero_grad()
+
+        outputs = self._net(inputs)
+        loss = self._criterion(outputs, labels)
+        loss.backward()
+        self._optimizer.step()
+
+        # Note that we save the value of the loss function to the _loss_items dictionary.
+        # This allows the model to automatically gather the losses for us if we are
+        # training in distributed mode.
+        self._loss_items["loss"] = loss
+
+    def on_training_batch_end(
+        self,
+        state: hlt.TrainingState,
+        should_log: bool = False,
+    ) -> None:
+        """
+        Perform steps after the training batch.
+
+        Args:
+            state (pyt.TrainingState): the current training state.
+            should_log (bool): if true, then writing to the log should be performed.
+        """
+        # If we were training in distributed mode, calling the base Model's function will
+        # automatically gather all loss values saved to self._loss_items.
+        super().on_training_batch_end(state, should_log)
+
+        # This flag is set to true whenever the number of iterations is a multiple of the
+        # logging frequency we set when creating the trainer.
+        if should_log:
+            root_logger = logging.get_root_logger()
+            tb_logger = hlc.get_from_optional(logging.get_tensorboard_writer())
+
+            loss_val = self._loss_items["loss"]
+
+            root_logger.info(
+                f"[{state.global_epoch + 1}, {state.global_iteration:5d}] "
+                f"loss: {loss_val:.3f}, "
+                f"running loss: {loss_val / state.running_iter:.3f} "
+                f"avg time: {state.average_iter_time:.2f}s"
+            )
+            tb_logger.add_scalar("train/loss", loss_val, state.global_iteration)
+            tb_logger.add_scalar(
+                "train/running loss",
+                loss_val / state.running_iter,
+                state.global_iteration,
+            )
+
+    def on_training_end(self) -> None:
+        """Perform steps after training ends."""
+        # For our example, we're going to save the hyper-params to the tensorboard log so
+        # we can compare with other runs.
+        # Notice that self._val_scores is active because we validate every epoch. If your
+        # validation frequency is different, you may need to alter this code.
+        total = self._val_scores["total"]
+        correct = self._val_scores["correct"]
+        accuracy = 100 * correct // total
+        writer = hlc.get_from_optional(logging.get_tensorboard_writer())
+        writer.add_hparams(
+            {"lr": 0.001, "momentum": 0.9, "epochs": 2},
+            {"hparam/accuracy": accuracy, "hparam/loss": self._loss_items["loss"].item()},
+        )
+
+    def eval(self) -> None:
+        """Set the model to eval."""
+        self._net.eval()
+
+    def on_validation_start(self, validation_cycle: int) -> None:
+        """Perform steps on validation start."""
+        # The base function will automatically clear the validation scores for us.
+        super().on_validation_start(validation_cycle)
+
+        # If you need to add further data to the table, you can do so here. In our case,
+        # we're going to add the total number of labels seen and how many of those were
+        # correct so we can compute the accuracy metric.
+        self._val_scores["total"] = 0
+        self._val_scores["correct"] = 0
+
+    def valid_step(self, batch: typing.Any, step: int) -> None:
+        """Perform the validation step."""
+        images, labels = batch
+        images = images.to(self.device)
+        labels = labels.to(self.device)
+
+        outputs = self._net(images)
+
+        _, predicted = torch.max(outputs.data, 1)
+        self._val_scores["total"] += labels.size(0)
+        self._val_scores["correct"] += (predicted == labels).sum().item()
+
+    def on_validation_end(self, validation_cycle: int) -> None:
+        """Perform steps after validation ends."""
+        root_logger = logging.get_root_logger()
+        tb_logger = hlc.get_from_optional(logging.get_tensorboard_writer())
+
+        # Grab the validation scores and compute the accuracy metric so we can log it. If
+        # we were in distributed mode, you would need to gather the values here.
+        total = self._val_scores["total"]
+        correct = self._val_scores["correct"]
+        accuracy = 100 * correct // total
+
+        root_logger.info(f"[Validation {validation_cycle}] accuracy: {accuracy}")
+        tb_logger.add_scalar("val", accuracy, validation_cycle)
+
+
+if __name__ == "__main__":
+    # Set the CUBLAS workspace setting to allow determinism to be used in CUDA >= 10.2.
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+    datamodule = CIFARDataModule(pathlib.Path.cwd())
+    model = ClassifierModel()
+
+    trainer = hlt.Trainer(
+        run_name="cifar10",
+        train_unit=hlt.TrainingUnit.EPOCH,
+        total_steps=2,
+        valid_frequency=1,
+        chkpt_frequency=1,
+        print_frequency=10,
+        enable_tensorboard=True,
+        enable_file_logging=True,
+        enable_progress_bar=True,
+        enable_deterministic=True,
+        chkpt_root=pathlib.Path.cwd() / "chkpt",
+        log_path=pathlib.Path.cwd() / "logs",
+        run_path=pathlib.Path.cwd() / "runs",
+    )
+
+    trainer.fit(model, datamodule)