micromind package

micromind.core module

Core class for micromind. Supports helper function for exports. Out-of-the-box multi-gpu and FP16 training with HF Accelerate and much more.

Authors:

• Francesco Paissan, 2023

class micromind.core.Metric(name: str, fn: Callable, reduction: str | None = 'mean', eval_only: bool | None = False, eval_period: int | None = 1)

Bases: object

Class for tracking evaluation metrics during training.

This class allows you to create custom evaluation metrics by providing a function to compute the metric and specifying a reduction method.

Parameters:

• name (str) – The name of the metric.

• fn (Callable) – A function that computes the metric given predictions and batch data.

• reduction (Optional[str]) – The reduction method for the metric (‘sum’ or ‘mean’). Default is ‘mean’.

Returns:

• Reduced metric. Optionally, you can access the metric history

• before call reduce(clear=True) (torch.Tensor)

Example

>>> from micromind import Metric, Stage
>>> import torch

>>> def custom_metric(pred, batch):
...     # Replace this with your custom metric calculation
...     return pred - batch

>>> metric = Metric("Custom Metric", custom_metric, reduction="mean")
>>> pred = torch.tensor([1.0, 2.0, 3.0])
>>> batch = torch.tensor([0.5, 1.5, 2.5])
>>> metric(pred, batch, stage=Stage.train)
>>> metric.history
{0: [tensor([0.5000, 0.5000, 0.5000])], 1: [], 2: []}
>>> metric.reduce(Stage.train)
0.5

reduce(stage, clear=False)

Compute and return the metric for a given prediction and batch data.

Parameters:

• pred (torch.Tensor) – The model’s prediction.

• batch (torch.Tensor) – The ground truth or target values.

• stage (Stage) – The current stage (e.g., Stage.train).

• Optional[str] (device) – The device on which to perform the computation. Default is ‘cpu’.

class micromind.core.MicroMind(hparams=None)

Bases: ABC

MicroMind is an abstract base class for creating and training deep learning models. Handles training on multi-gpu via accelerate (using DDP and other distributed training strategies). It automatically handles the device management for the training and the micromind’s export capabilities to onnx, OpenVino and TFLite.

Parameters:

hparams (Optional[Namespace]) – Hyperparameters for the model. Default is None.

abstract compute_loss(pred, batch)

Computes the cost function for the optimization process. It return a tensor on which backward() is called. This method should be overwritten for the specific application.

Parameters:

• pred (Union[torch.Tensor, Tuple]) – Output of the forward() function

• batch (torch.Tensor) – Batch as defined from the DataLoader.

Returns:

loss – Compute cost function.

Return type:

torch.Tensor

configure_optimizers()

Configures and defines the optimizer for the task. Defaults to adam with lr=0.001; It can be overwritten by either passing arguments from the command line, or by overwriting this entire method. Scheduler step is called every optimization step.

Returns:

: Union[Tuple[torch.optim.Adam, None], torch.optim.Adam]

Return type:

Optimizer and learning rate scheduler.

eval()

export(save_dir: Path | str, out_format: str = 'onnx', input_shape=None) → None

Export the model to a specified format for deployment. TFLite and OpenVINO need a Linux machine to be exported.

Parameters:

• save_dir (Union[Path, str]) – The directory where the exported model will be saved.

• out_format (Optional[str]) – The format for exporting the model. Default is ‘onnx’.

• input_shape (Optional[Tuple]) – The input shape of the model. If not provided, the input shape specified during model creation is used.

abstract forward(batch)

Forward step of the class. It gets called during inference and optimization. This method should be overwritten for specific applications.

Parameters:

batch (torch.Tensor) – Batch as output from the defined DataLoader.

Returns:

pred – Predictions - this depends on the task.

Return type:

Union[torch.Tensor, Tuple]

init_devices()

Initializes the data pipeline and modules for DDP and accelerated inference. To control the device selection, use accelerate config.

load_modules(checkpoint_path: Path | str)

Loads models for path.

Parameters:

checkpoint_path (Union[Path, str]) – Path to the checkpoint where the modules are stored.

on_train_end()

Runs at the end of each training. Cleans up before exiting.

on_train_start()

Initializes the optimizer, modules and puts the networks on the right devices. Optionally loads checkpoint if already present.

This function gets executed at the beginning of every training.

set_input_shape(input_shape: Tuple = (3, 224, 224))

Setter function for input_shape.

Parameters:

input_shape (Tuple) – Input shape of the forward step.

test(datasets: Dict = {}, metrics: List[Metric] = []) → None

Runs the test steps.

Parameters:

• datasets (Dict) – Dictionary with the test DataLoader. Should be present in the key test.

• metrics (List[Metric]) – List of metrics to compute during test step.

Returns:

Metrics computed on test set.

Return type:

Dict[torch.Tensor]

train(epochs: int = 1, datasets: Dict = {}, metrics: List[Metric] = [], checkpointer=None, debug: bool = False) → None

This method trains the model on the provided training dataset for the specified number of epochs. It tracks training metrics and can optionally perform validation during training, if the validation set is provided.

Parameters:

• epochs (int) – The number of training epochs.

• datasets (Dict) – A dictionary of dataset loaders. Dataloader should be mapped to keys “train”, “val”, and “test”.

• metrics (Optional[List[Metric]]) – A list of metrics to track during training. Default is an empty list.

• debug (bool) – Whether to run in debug mode. Default is False. If in debug mode, only runs for few epochs and with few batches.

validate() → Dict

Runs the validation step.

class micromind.core.Stage(train: int = 0, val: int = 1, test: int = 2)

Bases: object

enum to track training stage

test: int = 2

train: int = 0

val: int = 1

micromind.convert module

Conversion from pytorch to different standard formats for inference (ONNX, OpenVINO, tflite).

Authors:

• Francesco Paissan, 2023

• Alberto Ancilotto, 2023

micromind.convert.convert_to_onnx(net: Module | MicroMind, save_path: Path | str = 'model.onnx', simplify: bool = False, replace_forward: bool = False)

Converts nn.Module to onnx and saves it to save_path. Optionally simplifies it. This function is internally used from mm.MicroMind.

Parameters:

• net (Union[nn.Module, mm.MicroMind]) – PyTorch module to be exported.

• save_path (Union[Path, str]) – Output path for the ONNX model.

• simplify (bool) – True if you want to simplify the model. Defaults to False.

• replace_forward (bool) – Used if you want to replace the forward method. It is need if you are calling this function on a mm.MicroMind. Defaults to False.

Returns:

The path of the ONNX model.

Return type:

Path

micromind.convert.convert_to_openvino(net: Module | MicroMind, save_path: Path, replace_forward: bool = False) → str

Converts model to OpenVINO. Uses ONNX in the process and converts networks from channel-first to channel-last (for optimized inference).

Parameters:

• net (nn.Module) – PyTorch module to be exported.

• save_path (Union[Path, str]) – Output path for the OpenVINO model.

• replace_forward (bool) – Used if you want to replace the forward method. It is need if you are calling this function on a mm.MicroMind. Defaults to False.

Returns:

The path of the XML model.

Return type:

str

micromind.convert.convert_to_tflite(net: Module | MicroMind, save_path: Path | str, batch_quant: Tensor = None, replace_forward: bool = False) → None

Converts nn.Module to tf_lite, optionally quantizes it.

Parameters:

• net (nn.Module) – PyTorch module to be exported.

• save_path (Union[Path, str]) – Output path for the OpenVINO model.

• batch_quant (torch.Tensor) – Optional batch for quantization. When passed, it is used to create the statistics of the quantized activations.

• replace_forward (bool) – Used if you want to replace the forward method. It is need if you are calling this function on a mm.MicroMind. Defaults to False.

Subpackages

• micromind.networks package
  • Submodules
  • micromind.networks.phinet module
  • micromind.networks.xinet module
  • micromind.networks.yolo module
• micromind.utils package
  • Submodules
  • micromind.utils.checkpointer module
  • micromind.utils.helpers module
  • micromind.utils.yolo module