Metadata-Version: 2.4
Name: torch_utils
Version: 1.0.1
Summary: Library for Torch projects
Author: Corentin
Author-email: Corentin <corentin@ayo.tokyo>
License-Expression: BSD-3-Clause
License-File: LICENSE
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: POSIX :: Linux
Requires-Dist: numpy>=2.2.0
Requires-Dist: torch>=2.11
Requires-Dist: pytest ; extra == 'test'
Requires-Python: >=3.11
Project-URL: Homepage, https://codeberg.org/ayo/torch_utils
Project-URL: Issues, https://codeberg.org/ayo/torch_utils/issues
Provides-Extra: test
Description-Content-Type: text/markdown

# 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 the 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 multi-processed data loading, making the batch generator several times faster.
* `pipeline` is a function that 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).
* `save` parameter 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 useful 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 `BatchGenerator`. A `with` statement can be used to be safe. Upon deletion `BatchGenerator` object should also automatically call `release` but it is not safe as exception or other kind of interruption could bypass such code path.

A typical use should look like:

``` python
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 early 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_processor` function to be called to retrieve data and label each batch. The implementation will be added in latter versions.

### Sequence Batch Generator

Sequence 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:

``` python
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:


``` python
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:

``` python
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 mixed 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 writing 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 useful to quickly check if layers are over/under-sized. Also a `code` folder is created to save the current code base (`.py` files only) if using a [src layout](https://packaging.python.org/en/latest/discussions/src-layout-vs-flat-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 useful when the first loss/accuracy values are too high, making the summary plotting hard to read without giving useful 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 useful 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, useful 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:

``` python
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 useful 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:

``` sh
make dev
```

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