OverviewHardwareSoftwareCloud
OAK Apps
DepthAI
AI
Software Stack
Troubleshooting

ON THIS PAGE

  • Overview
  • Configuration
  • Configuration Components
  • Example Configuration File
  • Predefined Models
  • Customization
  • Custom Components
  • Example: Custom Callback
  • Example: Custom Loss Function
  • Using Custom Components in Configuration
  • Including Custom Components in Training

Concepts

Overview

Before diving into data preparation and model training, it's crucial to understand the core concepts that make LuxonisTrain highly flexible and customizable. This section will walk you through the fundamentals of Configuration and Customization.

Configuration

LuxonisTrain uses YAML configuration files to define every aspect of your training pipeline, from model architecture to optimization settings, offering flexibility in how you configure your training process. This page will guide you through the various configuration options available, ensuring you can tailor the tool to your needs. For a complete list of all parameters, their possible values, and default settings, please visit the LuxonisTrain configuration reference.

Configuration Components

  • Model: Defines the architecture of your model, including nodes, losses, visualizers and metrics.
  • Loader: Configures data loading by selecting the appropriate dataset. For more details on how to create a dataset, refer to the Data Preparation section.
  • Trainer: Specifies training parameters, including batch size, epochs, preprocessing (normalization, augmentations), optimizer, scheduler, callbacks, and more. For more details, refer to the Training section.
  • Tracker: Monitors and logs training metrics for analysis. For more details, refer to the Tracking section.
  • Exporter: Defines export, archiving, and conversion settings for trained models. For more details, refer to the Exporting section.
  • Tuner: Tunes hyperparameters to optimize model performance. For more details, refer to the Hyperparameter Tuning section.

Example Configuration File

Here's a sample YAML configuration file to give you an idea of how to structure your LuxonisTrain setup:
Yaml
1model:
2  name: model_name
3
4  # Use a predefined detection model instead of manually defining the architecture
5  predefined_model:
6    name: DetectionModel
7    params:
8      variant: light
9      loss_params:
10        iou_type: "siou"
11        n_warmup_epochs: 0
12
13# Dataset configuration: Download and parse the COCO dataset from Roboflow
14loader:
15  params:
16    dataset_name: coco_test
17    dataset_dir: "roboflow://team-roboflow/coco-128/2/coco"
18
19trainer:
20  batch_size: 8
21  epochs: 200
22  n_workers: 8
23  validation_interval: 10
24
25  preprocessing:
26    train_image_size: [384, 384]
27
28    # Image normalization using ImageNet standards
29    normalize:
30      active: true
31
32    # Image augmentations using Albumentations library
33    augmentations:
34      - name: Defocus
35      - name: Sharpen
36      - name: Flip
37
38  callbacks:
39    - name: ConvertOnTrainEnd
40    - name: TestOnTrainEnd
41
42  optimizer:
43    name: SGD
44    params:
45      lr: 0.02
46
47  scheduler:
48    name: ConstantLR
The dataset in the example above is parsed using the LuxonisML parser. To learn more about it, refer to the LuxonisParser documentation.This example demonstrates the use of a predefined detection model, which you can explore further here.

Predefined Models

You don't need to start designing architecture from scratch; you can use a predefined model based on the task you want to solve. Those models are well tested to be fast and accurate. You can check out all of them here.

Customization

LuxonisTrain's modular architecture allows you to customize and extend the framework to fit your specific needs. You can customize components like Loaders, Nodes, Losses, Metrics, Visualizers, Callbacks, Optimizers, and Schedulers. Additionally, the training strategy can be fully customized, allowing you to implement different learning rates for different parameter groups, apply various schedulers to different weights, and create complex warmup phases. This granular control enables sophisticated optimization strategies that can significantly improve model convergence and performance.

Custom Components

To implement a custom component, subclass the respective base class and register it. Here are some examples:

Example: Custom Callback

Python
1import lightning.pytorch as pl
2
3from luxonis_train import LuxonisLightningModule
4from luxonis_train.registry import CALLBACKS
5
6
7@CALLBACKS.register()
8class CustomCallback(pl.Callback):
9    def __init__(self, message: str, **kwargs):
10        super().__init__(**kwargs)
11        self.message = message
12
13    # Will be called at the end of each training epoch.
14    # Consult the PyTorch Lightning documentation for more callback methods.
15    def on_train_epoch_end(
16        self,
17        trainer: pl.Trainer,
18        pl_module: LuxonisLightningModule,
19    ) -> None:
20        print(self.message)

Example: Custom Loss Function

Python
1from torch import Tensor
2
3from luxonis_train import BaseLoss, Tasks
4
5
6# Subclasses of `BaseNode`, `BaseLoss`, `BaseMetric`
7# and `BaseVisualizer` are registered automatically.
8class CustomLoss(BaseLoss):
9    supported_tasks = [Tasks.CLASSIFICATION, Tasks.SEGMENTATION]
10
11    def __init__(self, smoothing: float, **kwargs):
12        super().__init__(**kwargs)
13        self.smoothing = smoothing
14
15    def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
16        # Implement the actual loss logic here.
17        value = predictions.sum() * self.smoothing
18        return value.abs()

Using Custom Components in Configuration

To use custom components in your training configuration, specify the component name and parameters in the YAML file. Here's an example:
Yaml
1model:
2  nodes:
3  - name: SegmentationHead
4    losses:
5    - name: CustomLoss
6      params:
7        smoothing: 0.0001
8
9trainer:
10  callbacks:
11    - name: CustomCallback
12      params:
13        lr: "Hello from the custom callback!"

Including Custom Components in Training

After creating custom components, you need to import them before starting training. There are two ways to accomplish this:

Using the CLI

Command Line
1luxonis_train --source custom_components.py train --config config.yaml

Using the Python API

Python
1# Import your custom components first
2from custom_components import *
3from luxonis_train import LuxonisModel
4
5model = LuxonisModel("config.yaml")
6model.train()
Either method ensures your custom components are registered and available during the training process.