Torch Utils

torch-utils is a library with tools to train neural networks with torch. Some tools such as BatchGenerator are very powerful and will accelerate trainings. Others like Trainer can be used as framework/boilerplate to avoid coding the same lines every time.

This library is python-only and can be easily modified or built upon.

Instalation

The library requires Python 3.11+. Then simply install from pip or equivalent:

pip install torch-utils

Usage

Batch Generator

The BatchGenerator is a object that can load data and labels then feed them during training as batches (set with batch_size). It keeps tracks of steps (global steps and epoch steps) and epochs (an epoch is a pass over all teh available data). It also deals with shuffling and it is made to be as efficient as possible.

The main issues with feeding data for a training session is that the dataset usually can fit in memory thus need to be loaded on the fly. BatchGenerator objects have data_processor and label_processor arguments to give them the function to transform raw data/label to training ones. For instance images can be fed as paths and a data_processor set to load (resize or any transformation) from path and return numpy arrays.

preload is an argument to run the data/label processor over the data and save everything to memory. If the dataset fits in memory this is the fastest way to feed data.

prefetch is a boolean argument to enable preparing the next batch in memory to be fed when next_batch is called. This is the fastest way to fed data during a training session if the dataset cannot fit in memory.

num_workers can be set to enable multiprocessed data loading, making the batch generator several times faster.

A pipeline function can be set to manipulate a batch (data+label) before it is given. Contrary to data/label processor this function doesn't not benefit from preload, prefetch or num_worker. This is used for data augmentation (random image flipping for instance) or any preparation that need to be run on the fly and is not deterministic. Because image flipping is very common flip_data is a boolean that will randomly flip vertically and/or horizontally like a pipeline (if a pipeline is set this flip is called first).

A save option is also available to give a path to save/load the processed data and label (after processed by the given data_processor/label_processor) on disk. This is usefull when the dataset fits in memory with preload but the processors take a lot of time thus saving the result on disk enable for near instant loading of the dataset. This is especially useful on early iteration when using a small portion of a dataset to develop a machine learning model as any code change can then be tried instantly*.

The release() function needs to be called before stopping a training session, this is necessary to clean processes and shared memory used if num_worker and/or prefecth is used. There are scenarios where this is not needed but it is a safer way of using BatchGenerato. A with statement can be used to be safe. Upon deletion BatchGenerator object should also automatically call release but it is not safe as expection or other kind of interuption could bypass such code path.

A typical use should look like:

import json
from pathlib import Path

import numpy as np
from PIL import Image

from torch_utils.batch import BatchGenerator


def data_processor(frame_path: Path) -> np.ndarray:
    image = Image.open(frame_path)
    return np.asarray(image, dtype=np.uint8).transpose([2, 0, 1])


data = data_dir.rglob('*.jpg')
label = json.loads(label_path.read_text(encoding='utf-8'))
max_epoch = 100
with BatchGenerator(
        data, label, data_processor=data_processor, batch_size=64, prefetch=True,
        num_workers=8) as batch_gen:
    import torch  # importing late so any ealy error can be caught quickly as torch is heavy to load

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    try:
        while batch_gen.epoch < max_epoch:
            batch_data = torch.as_tensor(batch_gen.batch_data, dtype=torch.uint8, device=device)
            batch_label = torch.as_tensor(batch_gen.batch_label, dtype=torch.float32, device=device)
            # Training with batch_data and batch_label
            print(f'Step {self.batch_train.global_step}', end='\r')
            self.batch_train.next_batch()
        print('\nTraining done')
    except KeyboardInterrupt:
        print('\rStopping training...')

IPC Batch Generator

IPC batch generator is similar to a batch generator expect that data and label are replaced by an ipc_processer function to be called to retrieve data and label each batch. The implementation will be added on a latter version.

Sequence Batch Generator

Sequance batch generators have an additionnal sequence_size parameter to set the length of desired sequences of data in the data/label output batch.

Layer

torch_utils layer module include several usual layers that inherit from global "configuration" from the Layer class set in _layer.py as:

class Layer(nn.Module):
    # Default layer arguments
    ACTIVATION: type[nn.Module] | None = nn.LeakyReLU
    ACTIVATION_KWARGS: ClassVar[dict] = {'negative_slope': 0.1}

    USE_BATCH_NORM = True
    BATCH_NORM_TRAINING = True
    BATCH_NORM_MOMENTUM = 0.01

This allow for easier swap of activation or batch normalization in a model. For instance a simple model can be defined as:

import torch
from torch import nn
from torch_utils.layer import Conv2d, Linear


class Model(nn.Module):
    def __init__(self, input_sample: torch.Tensor):
        super().__init__()
        self.input_shape = input_sample.shape[-3:]
        self.feature_extractor = nn.Sequential(
            Conv2d(3, 16, kernel_size=5, stride=4, padding=2),
            Conv2d(16, 16, stride=2, padding=1),
            Conv2d(16, 8, stride=2, padding=1),
            Conv2d(8, 8, stride=2, padding=1),
            Conv2d(8, 8, stride=2, padding=1),
            Conv2d(8, 8, stride=2, padding=1),
            Conv2d(8, 8, stride=2, padding=1))

        out_shape = self.feature_extractor(input_sample.cpu()).shape
        self.flatten_shape = out_shape[-1] * out_shape[-2] * out_shape[-3]
        self.detector = nn.Sequential(
            Linear(self.flatten_shape, 64, use_batch_norm=False),
            Linear(64, 64, use_batch_norm=False),
            Linear(64, 32, activation=None, use_batch_norm=False),
            Linear(32, 4, activation=nn.Sigmoid, use_batch_norm=False))

    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
        return self.detector(self.feature_extractor(input_data).view(-1, self.flatten_shape))

By changing the Layer class attribute default parameters can be easily swaped:

import torch
from torch_utils.layer import Layer
from .network import Model

def train():
    batch_data: torch.Tensor

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    Layer.BATCH_NORM_MOMENTUM = 1e-3  # can swap batch norm mometum for the whole network
    network = Network(batch_data).to(device)

    out = model(batch_data)

Layers are merely wrapper around torch.nn layers so both can be mized together without any issue.

Current available layers are: Linear, Conv1d, Conv2d, Conv3d, Deconv2d.

Trainer

Trainer is a class used to train a model using its fit methods. This avoid writting boilerplate code to train a model. Looking at the fit method's code will allow anyone to copy and customized their model training.

Trainer objects will create tensorboard summaries with loss, accuracy (if provided) and learning rate (if scheduler is provided) information automatically. This is done in the save_summaries method. A summary_callback can be overloaded to add additional summaries.

Mandatory parameters are:

• batch_generator_train/batch_generator_val: the first 2 parameters are BatchGenerator object used to get train/validation data.
• device: Trainer is using torch.as_tensor from data given by the batch generators, this argument is passed as is. This should be the same device as the given network.
• output_dir: path where to save models and tensorboard metrics (under tensorboard). Trainer is also saving a parameters.csv with all the network parameters info (shape and sizes), this is usefull to quickly check if layers are over/under-sized. Also a code folder is created to save the current code base (.py files onlt) if using a src layout. And finally the initial (not yet trained) model is saved as model_init.pt.
• network: model to be trained.
• optimizer: optimizer to be used.
• loss: can be a single torch Module or any mapping (key-value like dict or list of tuple) of name (as string) and modules. Giving a map will allow the Trainer to compute the losses separately to be included in metrics before combined for back-propagation.

Optional arguments are:

• summary_per_epoch: number of summaries to be saved per epoch, minimum accepted being 1.
• epoch_skip : avoid creating summaries for the first given epochs. This is usefull when the first loss/accuracy values are too high, making the summary plotting hard to read without giving usefull information.
• data_dtype/label_dtype: Trainer is using torch.as_tensor from data given by the batch generators, those arguments are past to cast to any desired types. Default is set to None (no casting).
• accuracy: torch Module used for computing accuracy, will automatically be added to the saved metrics.
• scheduler: torch LRSchedule (Learning Rate Scheduler) to be used.
• pre_process: torch Module to be used on a batch before being fed to the model. This is usefull for data normalization as BatchGenerator is meant to run using numpy arrays thus on CPU while some data preparation can be faster on GPU but the pre-process is not meant to be saved with the model.
• train_pre_process: same as pre_process but will be used instead during training, making pre_process (if defined) only used in validation steps.
• data_is_label: boolean used when training model such as auto-encoders where the loss should take the input data itself as label.
• max_checkpoints: limits the number of saved models (checkpoint), only the "best" models are kept. By default best models are chosen from having lowest validation loss, if accurary is set lowest validation accuracy is used.
• logger: passing logger to be used. Trainer is sending 3 different info-level messages: at the begining of a fit call (epoch, steps per epoch, etc), at the end of a fit call (training time, memory and VRAM usage) and after each metric save (current epoch, step, loss, and training speed)
• verbose: Print progress during training, with losses information and speed. Default is True.
• save_src: Saving code (only .py files) used during training (work only for src layout). Default is True.

Multiple "callback" methods can be overloaded for customization:

• train_step_callback and val_step_callback at the end of train/validation steps.
• end_epoch_callback at the end of training epochs.
• pre_summary_callback before creating summaries (metrics), also before validation loop used during summary.
• summary_callback at the end of summary saving, usefull to add additional metrics.

Each callback has specific sets of parameters. Trainer creates automatically writer_train and writer_val that can be used in callbacks from self.

A typical usage of Trainer would look like:

from datetime import datetime
import logging

import torch
from torch_utils.batch import BatchGenerator
from torch_utils.trainer import Trainer
from torch_utils.utils.logger import create_logger


class TrainConfig:
    BATCH_SIZE = 64
    TRAIN_EPOCH = 100
    LEARNING_RATE = 1e-3
    SCHEDULER_GAMMA = 1 - 5e-4

    MAX_CHECKPOINT = 50
    MAX_IMAGES_OUTPUT = 4
    SKIP_FIRST_EPOCHS = 10
    SUMMARIES_PER_EPOCH = 1


logger = create_logger('trainer', output_dir / 'log', stdout=True)
output_dir = Path('data', 'train', datetime.today().strftime('%Y-%m-%d %H:%M:%S'))
with BatchGenerator(...) as batch_train, BatchGenerator(...) as batch_val:
    network = MyNetwork().to(device)
    optimizer = torch.optim.Adam(network.parameters(), lr=TrainConfig.LEARNING_RATE)
    loss = torch.nn.MSELoss()
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=TrainConfig.SCHEDULER_GAMMA)

    class MyTrainer(Trainer):
        @override
        def summary_callback(
                self, train_inputs: torch.Tensor, train_processed: torch.Tensor,
                train_labels: torch.Tensor, train_outputs: torch.Tensor, train_running_count: int,
                val_inputs: torch.Tensor, val_processed: torch.Tensor,
                val_labels: torch.Tensor, val_outputs: torch.Tensor, val_running_count: int):
            nonlocal batch_train, image_shape

            for writer, inputs, labels, outputs in [
                    (self.writer_train, train_processed, train_labels.detach(), train_outputs.detach()),
                    (self.writer_val, val_processed, val_labels.detach(), val_outputs.detach())]:
                for input_index in range(min(inputs.shape[0], TrainConfig.MAX_IMAGES_OUTPUT)):
                    # typical de-normalization
                    image = (inputs[input_index] * 255.).type(torch.uint8).permute((1, 2, 0))
                    # Draw on image from labels/outputs
                    writer.add_image(
                        f'image/{input_index}', image, global_step=batch_train.global_step, dataformats='HWC')

    trainer = MyTrainer(
        batch_train, batch_val, device=device, output_dir=output_dir, network=network, optimizer=optimizer,
        loss=loss, epoch_skip=TrainConfig.SKIP_FIRST_EPOCHS, summary_per_epoch=TrainConfig.SUMMARIES_PER_EPOCH,
        accuracy=accuracy, scheduler=scheduler, max_checkpoints=TrainConfig.MAX_CHECKPOINT, logger=logger)
    trainer.fit(TrainConfig.TRAIN_EPOCH)

Utils

utils contains helpers mostly meant as usage within torch_utils but still useful overall. create_logger from utils.logger is the only method not used internally but meant as helper for projects in general.

• logger: create_logger from this module will create a python's Logger, thus requires a name argument. Additional arguments are:
  • log_dir: Path to save logs (file path, default is set to None), if not set no file are saved. RotatingFileHandler is automatically used to rotate logs in case of large amount of data with a 500k limit and 5 files max. This ensures logging can be done without any limitation while never compromising storage space. Limits are hard-coded for now but could be customized in future version.
  • stdout: boolean to print logs in stdout, default is False.
  • memory: integer to add a logging buffer to write into. Can be usefull to inspect previous logs during for automatic testing, if not using the pytest log capture.
  • color: add color to stdout logs, this is using escape sequences only working on UNIX system (non Windows).
  • level: standard python's log level, can be used to avoid printing debug/info/warning logs. Default is logging.INFO (native logging module).
• memory: helper function related to memory number parsing.
  • human_size: from integer to text with Ki/Mi/Gi prefixes per 1024 power
  • human_to_bytes: reverse of human_size, parsing text back to number (as float by default).
• train: various train-related helper function.
  • parameter_summary: returns list of parameter information (name, shape, size) from given network (torch nn.Module)
  • ressource_usage: returns used memory and VRAM
  • gpu_usde_memory: returns used VRAM (nvidia only)
  • gpu_total_memory: returns total VRAM (nvidia only)

Other

transformer and ssd contains working code for well-known models that are not natively included in PyTorch. Documentation will be added if/when current implementation are sufficiently validated and tested.

Contributing

To start working on this repository, if you have uv installed you can simply:

make dev

After any change you can use make lint, make test and make coverage.