Skip to content
BranchesTags

Repository files navigation

Luxonis Training Framework

Ubuntu Windows MacOS

License PyBadge Ruff CI codecov

🌟 Overview

LuxonisTrain is a user-friendly tool designed to streamline the training of deep learning models, especially for edge devices. Built on top of PyTorch Lightning, it simplifies the process of training, testing, and exporting models with minimal coding required.

Title Image

✨ Key Features

  • No Coding Required: Define your training pipeline entirely through a single YAML configuration file.
  • Predefined Configurations: Utilize ready-made configs for common computer vision tasks to start quickly.
  • Customizable: Extend functionality with custom components using an intuitive Python API.
  • Edge Optimized: Focus on models optimized for deployment on edge devices with limited compute resources.

Warning

The project is in a beta state and might be unstable or contain bugs - please report any feedback.

🚀 Quick Start

Get started with LuxonisTrain in just a few steps:

  1. Install LuxonisTrain

    pip install luxonis-train

    This will create the luxonis_train executable in your PATH.

  2. Pick a bundled predefined model

    Every predefined model ships with the wheel. List them with:

    luxonis_train list-models

  3. Find a suitable dataset for your task

    We will use a sample COCO dataset from RoboFlow in this example.

    [!IMPORTANT] A roboflow:// source needs the ROBOFLOW_API_KEY environment variable. Get your key from the Roboflow settings. The Roboflow documentation gives the steps.

  4. Start training

    luxonis_train train                  \
  --model detection --variant light  \
  loader.params.dataset_dir "roboflow://team-roboflow/coco-128/2/coco"

    --model detection --variant light runs the bundled detection_light_model.yaml configuration. Everything after the options is an override applied on top of it, which is how you point the configuration at your data. Once you have your own dataset registered, the usual form is:

    luxonis_train train                  \
  --model detection --variant light  \
  loader.params.dataset_name "my_dataset"

    To change more than a handful of fields, write your own config file and pass it with --config instead. See Configuration.

  5. Monitor progress with TensorBoard

    tensorboard --logdir output/tensorboard_logs

    Open the provided URL in your browser to visualize the training progress

Note

For hands-on examples of how to prepare data with LuxonisML and train AI models using LuxonisTrain, check out this guide.

📜 Table Of Contents

🛠️ Installation

LuxonisTrain requires Python 3.10 or higher. We recommend using a virtual environment to manage dependencies.

Install via pip:

pip install luxonis-train

This will also install the luxonis_train CLI. For more information on how to use it, see CLI Usage.

AIMET Quantization Support

To enable support for AIMET quantization, install the luxonis-train[aimet] extra:

pip install luxonis-train[aimet]

📝 Usage

You can use LuxonisTrain either from the command line or via the Python API. We will demonstrate both ways in the following sections.

💻 CLI

The CLI is the most straightforward way how to use LuxonisTrain. The CLI provides several commands for training, testing, tuning, exporting and more.

Available commands:

  • train - Start the training process
  • test - Test the model on a specific dataset view
  • infer - Run inference on a dataset, image directory, or a video file.
  • export - Export the model to ONNX
  • archive - Create an NN Archive (.tar.xz) from an ONNX model
  • convert - Export + archive + platform-specific conversion (RVC2/RVC3/RVC4)
  • tune - Tune the hyperparameters of the model for better performance
  • inspect - Inspect the dataset you are using and visualize the annotations
  • annotate - Annotate a directory using the model’s predictions and generate a new LDF.
  • quantize - Quantize the model using AIMET quantization techniques
  • list-models - List packaged predefined models along with their variants and available versions

To get help on any command:

luxonis_train <command> --help

Specific usage examples can be found in the respective sections below.

Selecting a model. Every command that accepts --config <path> also accepts --model <name> (optionally with --variant). The two are mutually exclusive and answer different questions:

  • --config my_config.yaml runs your config file.
  • --model detection --variant light runs the matching packaged configuration (here detection_light_model.yaml), read straight from the installed package. There is no local file to download or edit, and a config file in your working directory is not picked up.

Both forms accept the same key value overrides, so you can adapt a packaged configuration from the command line:

# Train the packaged detection model on your own dataset
luxonis_train train                  \
  --model detection --variant light  \
  loader.params.dataset_name "my_dataset"

# Same, but pin the predefined-model version explicitly
luxonis_train train --model detection:v1 --variant light

Use luxonis_train list-models to see the available (model, variant) combos and version numbers, and luxonis_train info --model detection to see what a predefined model contains.

The Python API accepts the same packaged model selection directly:

from luxonis_train import LuxonisModel

model = LuxonisModel(
    model="detection:v1",
    variant="medium",
    opts={"loader.params.dataset_name": "my_dataset"},
)

Packaged configurations are a starting point, not a customization mechanism: as soon as you need to change more than a few fields (augmentations, losses, the training schedule), write your own config file. A few lines are enough, since it can build on the same predefined model:

model:
  predefined_model:
    name: DetectionModel
    params:
      variant: light

loader:
  params:
    dataset_name: my_dataset

trainer:
  epochs: 300

and run it with luxonis_train train --config my_config.yaml. The packaged YAMLs are good templates for this; they are visible in this repository and, in an installed environment, under the directory printed by python -c "from luxonis_train.config.predefined import configs_dir; print(configs_dir())".

Predefined-model versioning. The model.predefined_model block accepts an optional version field (default "latest"). When we ship a breaking architecture change to, say, DetectionModel, it will land in a new v2 package alongside the current one (predefined_models/detection/v2/); the class keeps its name, its version is inferred from the package, and it registers as DetectionModel:v2. Existing configs pinned to version: 1 keep resolving to the old architecture, and loading a checkpoint whose training-time version differs from the current-config version prints a clear warning telling you which version to pin. The equivalent CLI form is --model detection:v1, --model detection:v2, or --model detection:latest.

Note

CLI commands train, test, and tune can be run with --debug flag which allows the model to be used without a functional dataset.

⚙️ Configuration

LuxonisTrain uses YAML configuration files to define the training pipeline. Here's a breakdown of the key sections:

model:
  name: model_name

  # Use a predefined detection model instead of defining
  # the model architecture manually
  predefined_model:
    name: DetectionModel
    # Optional: pin an explicit architecture version. Defaults to
    # "latest". Pin to an integer (e.g. `version: 1`) to reproduce an
    # older checkpoint after a breaking change.
    version: latest
    params:
      variant: light

# Download and parse the coco dataset from RoboFlow.
# Save it internally as `coco_test` dataset for future reference.
loader:
  params:
    dataset_name: coco_test
    dataset_dir: "roboflow://team-roboflow/coco-128/2/coco"

trainer:
  batch_size: 8
  epochs: 200
  n_workers: 8
  validation_interval: 10

  preprocessing:
    train_image_size: [384, 384]

    # Uses the imagenet normalization by default
    normalize:
      active: true

    # Augmentations are powered by Albumentations
    augmentations:
      - name: Defocus
      - name: Sharpen
      - name: Flip

  callbacks:
    - name: ConvertOnTrainEnd
    - name: TestOnTrainEnd

  optimizer:
    name: SGD
    params:
      lr: 0.02

  scheduler:
    name: ConstantLR

📚 Configuration Reference

For a complete reference of all available configuration options, see our Configuration Documentation.

Tip

We provide a set of predefined configuration files for common computer vision tasks in the luxonis_train/configs directory. They ship with the package and are what --model / --variant select, so they can be used without a local copy. They are also great starting points to copy and customize for your specific needs.

🗃️ Data Preparation

LuxonisTrain supports several ways of loading data:

  • using a data directory in one of the supported formats
  • using an already existing dataset in our custom LuxonisDataset format
  • using a custom loader
    • to learn how to implement and use custom loaders, see Customizations

📂 Data Directory

The easiest way to load data is to use a directory with the dataset in one of the supported formats.

Supported formats:

  • COCO - We support COCO JSON format in two variants:
  • Pascal VOC XML
  • YOLO Darknet TXT
  • YOLOv4 PyTorch TXT
  • MT YOLOv6
  • CreateML JSON
  • TensorFlow Object Detection CSV
  • Classification Directory - A directory with subdirectories for each class 
    dataset_dir/
├── train/
│   ├── class1/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
│   │   └── ...
│   ├── class2/
│   └── ...
├── valid/
└── test/
  • Segmentation Mask Directory - A directory with images and corresponding masks. 
    dataset_dir/
├── train/
│   ├── img1.jpg
│   ├── img1_mask.png
│   ├── ...
│   └── _classes.csv
├── valid/
└── test/
    The masks are stored as grayscale PNG images where each pixel value corresponds to a class. The mapping from pixel values to classes is defined in the _classes.csv file. 
    Pixel Value, Class
0, background
1, class1
2, class2
3, class3

Preparing your Data

  1. Organize your dataset into one of the supported formats.
  2. Place your dataset in a directory accessible by the training script.
  3. Update the dataset_dir parameter in the configuration file to point to the dataset directory.

The dataset_dir can be one of the following:

  • Local path to the dataset directory
  • URL to a remote dataset
    • The dataset will be downloaded to a "data" directory in the current working directory
    • Supported URL protocols:
      • s3://bucket/path/to/directory fo AWS S3
      • gs://buclet/path/to/directory for Google Cloud Storage
      • roboflow://{workspace}/{project}/{version}/{format} for RoboFlow
        • workspace - name of the workspace the dataset belongs to
        • project - name of the project the dataset belongs to
        • version - version of the dataset
        • format - one of coco, darknet, voc, yolov4pytorch, mt-yolov6, createml, tensorflow, folder, png-mask-semantic
        • example: roboflow://team-roboflow/coco-128/2/coco
      • ultralytics://{username}/datasets/{slug} for Ultralytics
        • username - name of the dataset author
        • slug - name of the dataset
        • Optional ?v={version} attached at the end for specific version of the dataset
        • example: ultralytics://ultralytics/datasets/coco8

Example:

loader:
  params:
    dataset_name: "coco_test"
    dataset_dir: "roboflow://team-roboflow/coco-128/2/coco"

💾 LuxonisDataset

LuxonisDataset is our custom dataset format designed for easy and efficient dataset management. To learn more about how to create a dataset in this format from scratch, see the Luxonis ML repository.

To use the LuxonisDataset as a source of the data, specify the following in the config file:

loader:
  params:
    # name of the dataset
    dataset_name: "dataset_name"

    # one of local (default), s3, gcs
    bucket_storage: "local"

Tip

To inspect the loader output, use the luxonis_train inspect command:

luxonis_train inspect --model detection --variant light

--config my_config.yaml works the same, with the loader section of your own config.

The inspect command is currently only available in the CLI

For additional information about the shapes of Luxonis ML data that the loader returns, please refer to the Loaders README.

🏋️‍♂️ Training

Once your configuration file and dataset are ready, start the training process.

CLI:

# your own config file
luxonis_train train --config my_config.yaml

# or a packaged predefined model, unmodified
luxonis_train train --model detection --variant light

The --model form trains the packaged configuration as it ships, which is rarely what you want on its own. At the very least you need to provide your own dataset. Overrides do that without a config file of your own:

luxonis_train train                  \
  --model detection --variant light  \
  loader.params.dataset_name "my_dataset"

Tip

Any configuration parameter can be changed this way, whether the config came from --config or --model. The value is looked up by its dotted path in the config:

luxonis_train train \
  --model detection --variant light \
  loader.params.dataset_dir "roboflow://team-roboflow/coco-128/2/coco" \
  trainer.epochs 300 \
  trainer.batch_size 8

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel(
    "my_config.yaml",
    {
        "loader.params.dataset_dir": "roboflow://team-roboflow/coco-128/2/coco",
        "trainer": {
            "epochs": 300,
            "batch_size": 8,
        },
    },
)
model.train()

Expected Output:

INFO     Using predefined model: `DetectionModel`
INFO     Main metric: `MeanAveragePrecision`
INFO     GPU available: True (cuda), used: True
INFO     TPU available: False, using: 0 TPU cores
INFO     HPU available: False, using: 0 HPUs
...
INFO     Training finished
INFO     Checkpoints saved in: output/1-coral-wren

Monitoring with TensorBoard:

If not explicitly disabled, the training process will be monitored by TensorBoard. To start the TensorBoard server, run:

tensorboard --logdir output/tensorboard_logs

Open the provided URL to visualize training metrics.

✍ Testing

Evaluate your trained model on a specific dataset view (train, val, or test).

CLI:

luxonis_train test --model detection --variant light \
                   --view val                        \
                   --weights path/to/checkpoint.ckpt

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_config.yaml")
model.test(weights="path/to/checkpoint.ckpt")

If you plan to continue using the same LuxonisModel run after testing, for example to call export() or archive(), use model.test(..., finalize_tracker=False) to keep the tracker run open and call model.finalize_run() once the full sequence is complete.

The testing process can be started automatically at the end of the training by using the TestOnTrainEnd callback. To learn more about callbacks, see Callbacks.

🧠 Inference

Run inference on images, datasets, or videos.

CLI:

  • Inference on a Dataset View:
luxonis_train infer --model detection --variant light \
                    --view val                        \
                    --weights path/to/checkpoint.ckpt
  • Inference on a Video File:
luxonis_train infer --model detection --variant light \
                    --weights path/to/checkpoint.ckpt \
                    --source-path path/to/video.mp4
  • Inference on an Image Directory:
luxonis_train infer --model detection --variant light \
                    --weights path/to/checkpoint.ckpt \
                    --source-path path/to/images      \
                    --save-dir path/to/save_directory

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_config.yaml")

# infer on a dataset view
model.infer(weights="path/to/checkpoint.ckpt", view="val")

# infer on a video file
model.infer(weights="path/to/checkpoint.ckpt", source_path="path/to/video.mp4")

# infer on an image directory and save the results
model.infer(
    weights="path/to/checkpoint.ckpt",
    source_path="path/to/images",
    save_dir="path/to/save_directory",
)

🤖 Exporting

Export your trained models to ONNX for downstream conversion and deployment.

To configure the exporter, you can specify the exporter section in the config file. Note that exporter.hubai and exporter.blobconverter are only used by convert (or ConvertOnTrainEnd), not by export alone.

You can see an example export configuration here.

CLI:

luxonis_train export --config my_export_config.yaml --weights path/to/weights.ckpt

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_export_config.yaml")
model.export(weights="path/to/weights.ckpt")

Model export can be run automatically at the end of the training by using the ExportOnTrainEnd callback. For a full export + archive + platform conversion pipeline, use ConvertOnTrainEnd.

The exported models are saved in the export directory within your output folder.

🗂️ NN Archive

Create an NN Archive (.tar.xz) file for easy deployment with the DepthAI API. If you do not provide an ONNX executable, the model is exported to ONNX first.

The archive contains the exported model together with all the metadata needed for running the model.

CLI:

luxonis_train archive                \
  --model detection --variant light  \
  --weights path/to/checkpoint.ckpt

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_config.yaml")
model.archive(weights="path/to/checkpoint.ckpt")

The archive can be created automatically at the end of the training by using the ArchiveOnTrainEnd callback.

🔁 Convert

Convert is the unified flow for deployment. It performs:

  1. Export: .pt/.ckpt -> .onnx
  2. Archive: .onnx -> .tar.xz (NN Archive)
  3. Platform-specific conversion (optional): NN Archive -> platform NN Archive via HubAI SDK (recommended) or blobconverter (deprecated, RVC2 legacy .blob)

Configure conversion via the exporter section (exporter.hubai or exporter.blobconverter).

CLI:

luxonis_train convert --model detection --variant light --weights path/to/checkpoint.ckpt

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_config.yaml")
archive_path, conversion_artifacts = model.convert(
    weights="path/to/checkpoint.ckpt"
)

convert() returns a tuple where:

  • first item is the ONNX-based NN Archive path (.tar.xz)
  • second item is a dict[str, Path] with additional conversion artifacts (for example blob or hubai_archive)

Convert can be run automatically at the end of the training by using the ConvertOnTrainEnd callback.

🔬 Tuning

Optimize your model's performance using hyperparameter tuning powered by Optuna.

Configuration:

Include a tuner section in your configuration file. A full example is available here.

tuner:
  study_name: det_study
  n_trials: 10
  storage:
    backend: sqlite
  params:
    trainer.optimizer.name_categorical: ["Adam", "SGD"]
    trainer.optimizer.params.lr_float: [0.0001, 0.001]
    trainer.batch_size_int: [4, 16, 4]

CLI:

luxonis_train tune --config my_tuning_config.yaml

Python API:

from luxonis_train import LuxonisModel

model = LuxonisModel("my_tuning_config.yaml")
model.tune()

🎨 Customizations

LuxonisTrain is highly modular, allowing you to customize various components:

  • Loaders: Handles data loading and preprocessing.
  • Nodes: Represents computational units in the model architecture.
  • Losses: Define the loss functions used to train the model.
  • Metrics: Measure the model's performance during training.
  • Visualizers: Visualize the model's predictions during training.
  • Callbacks: Allow custom code to be executed at different stages of training.
  • Optimizers: Control how the model's weights are updated.
  • Schedulers: Adjust the learning rate during training.
  • Training Strategy: Specify a custom combination of optimizer and scheduler to tailor the training process for specific use cases.

Creating Custom Components:

Implement custom components by subclassing the respective base classes and/or registering them. Registered components can be referenced in the config file. Custom components need to inherit from their respective base classes:

Examples:

Custom Callback:

import lightning.pytorch as pl

from luxonis_train import LuxonisLightningModule
from luxonis_train.registry import CALLBACKS


@CALLBACKS.register()
class CustomCallback(pl.Callback):
    def __init__(self, message: str, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    # Will be called at the end of each training epoch.
    # Consult the PyTorch Lightning documentation for more callback methods.
    def on_train_epoch_end(
        self,
        trainer: pl.Trainer,
        pl_module: LuxonisLightningModule,
    ) -> None:
        print(self.message)

Custom Loss:

from torch import Tensor

from luxonis_train import BaseLoss, Tasks

# Subclasses of `BaseNode`, `BaseLoss`, `BaseMetric`
# and `BaseVisualizer` are registered automatically.
class CustomLoss(BaseLoss):
    supported_tasks = [Tasks.CLASSIFICATION, Tasks.SEGMENTATION]

    def __init__(self, smoothing: float, **kwargs):
        super().__init__(**kwargs)
        self.smoothing = smoothing

    def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
        # Implement the actual loss logic here
        value = predictions.sum() * self.smoothing
        return value.abs()

For additional examples of creating custom components, please refer to the examples section.

Using custom components in the configuration file:

model:
  nodes:
  - name: SegmentationHead
    losses:
    - name: CustomLoss
      params:
        smoothing: 0.0001

trainer:
  callbacks:
    - name: CustomCallback
      params:
        lr: "Hello from the custom callback!"

Note

Files containing the custom components must be sourced before the training script is run. To do that in CLI, you can use the --source argument:

luxonis_train --source custom_components.py train --config config.yaml

Python API:

You have to import the custom components before creating the LuxonisModel instance.

from custom_components import *
from luxonis_train import LuxonisModel

model = LuxonisModel("config.yaml")
model.train()

For more information on how to define custom components, consult the respective in-source documentation.

📚 Tutorials and Examples

We are actively working on providing examples and tutorials for different parts of the library which will help you to start more easily. The tutorials can be found here and will be updated regularly.

🔑 Credentials

When using cloud services, avoid hard-coding credentials or placing them directly in your configuration files. Instead:

  • Use environment variables to store sensitive information.
  • Use a .env file and load it securely, ensuring it's excluded from version control.

Supported Cloud Services:

  • AWS S3, requires:
    • AWS_ACCESS_KEY_ID
    • AWS_SECRET_ACCESS_KEY
    • AWS_S3_ENDPOINT_URL
  • Google Cloud Storage, requires:
    • GOOGLE_APPLICATION_CREDENTIALS
  • RoboFlow, requires:
    • ROBOFLOW_API_KEY

For logging and tracking, we support:

  • MLFlow, requires:
    • MLFLOW_S3_BUCKET
    • MLFLOW_S3_ENDPOINT_URL
    • MLFLOW_TRACKING_URI
  • WandB, requires:
    • WANDB_API_KEY

For remote database storage, we support:

  • POSTGRES_PASSWORD
  • POSTGRES_HOST
  • POSTGRES_PORT
  • POSTGRES_DB

🤝 Contributing

We welcome contributions! Please read our Contribution Guide to get started. Whether it's reporting bugs, improving documentation, or adding new features, your help is appreciated.

About

Training framework for easy creation, training and exporting of deep learning models.

Resources

Readme
Apache-2.0 license

Contributing

Contributing
Activity
Custom properties

Stars

27 stars

Watchers

4 watching

Forks

5 forks
Report repository

Releases

Packages

Used by

Contributors

Languages