TorchABC

torchabc is a lightweight package that provides an Abstract Base Class (ABC) to structure PyTorch projects and keep code well organized.

The core of the package is the TorchABC class. This class defines the abstract training and inference workflows and must be subclassed to implement a concrete logic.

This package has no extra dependencies beyond PyTorch and it consists of a simple self-contained file. It is ideal for research, prototyping, and teaching.

Structure

The TorchABC class structures a project into the following main steps:

1. Dataloaders - load raw data samples.
2. Preprocess – transform raw samples.
3. Collate - batch preprocessed samples.
4. Network - compute model outputs.
5. Loss - compute error against targets.
6. Optimizer - update model parameters.
7. Postprocess - transform outputs into predictions.

Each step corresponds to an abstract method in TorchABC. To use TorchABC, create a concrete subclass and implement these methods.

Quick start

Install the package.

pip install torchabc

Generate a template using the command line interface.

torchabc --create template.py --min

Fill out the template by implementing the methods below. The documentation of each method is available here.

import torch
from torchabc import TorchABC
from functools import cached_property


class MyModel(TorchABC):
    
    @cached_property
    def dataloaders(self):
        raise NotImplementedError
    
    @staticmethod
    def preprocess(sample, hparams, flag=''):
        return sample

    @staticmethod
    def collate(samples):
        return torch.utils.data.default_collate(samples)

    @cached_property
    def network(self):
        raise NotImplementedError
    
    @staticmethod
    def loss(outputs, targets, hparams):
        raise NotImplementedError

    @cached_property
    def optimizer(self):
        raise NotImplementedError
    
    @staticmethod
    def postprocess(outputs, hparams):
        return outputs

Usage

Once a subclass of TorchABC is implemented, it can be used for training, evaluation, checkpointing, and inference.

Initialization

model = MyModel()

Initialize the model.

Training

model.train(epochs=5, on="train", val="val")

Train the model for 5 epochs using the train and val dataloaders.

Evaluation

metrics = model.eval(on="test")

Evaluate on the test dataloader and return metrics.

Checkpoints

model.save("checkpoint.pth")
model.load("checkpoint.pth")

Save and restore the model state.

Inference

preds = model(samples)

Run predictions on raw input samples.

API Reference

The TorchABC class defines a standard workflow for PyTorch projects. Some methods are abstract (must be implemented in subclasses), others are optional (can be overridden but have defaults), and a few are concrete (should not be overridden).

Abstract Methods

Default Methods

Concrete Methods

Examples

Get started with simple self-contained examples:

• MNIST classification

Run the examples

Install the dependencies

poetry install --with examples

Run the examples by replacing <name> with one of the filenames in the examples folder

poetry run python examples/<name>.py

Contribute

Contributions are welcome! Submit pull requests with new examples or improvements to the core TorchABC class itself.