# 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](https://github.com/luxonis/luxonis-train/blob/main/configs/README.md).

### 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](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/data-preparation.md)
   section.
 * Trainer: Specifies training parameters, including batch size, epochs, preprocessing (normalization, augmentations), optimizer,
   scheduler, callbacks, and more. For more details, refer to the
   [Training](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/training.md) section.
 * Tracker: Monitors and logs training metrics for analysis. For more details, refer to the
   [Tracking](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/training.md) section.
 * Exporter: Defines export, archiving, and conversion settings for trained models. For more details, refer to the
   [Exporting](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/exporting.md) section.
 * Tuner: Tunes hyperparameters to optimize model performance. For more details, refer to the [Hyperparameter
   Tuning](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/training.md) section.

### Example Configuration File

Here's a sample YAML configuration file to give you an idea of how to structure your LuxonisTrain setup:

```yaml
model:
  name: model_name

  # Use a predefined detection model instead of manually defining the architecture
  predefined_model:
    name: DetectionModel
    params:
      variant: light
      loss_params:
        iou_type: "siou"
        n_warmup_epochs: 0

# Dataset configuration: Download and parse the COCO dataset from Roboflow
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]

    # Image normalization using ImageNet standards
    normalize:
      active: true

    # Image augmentations using Albumentations library
    augmentations:
      - name: Defocus
      - name: Sharpen
      - name: Flip

  callbacks:
    - name: ConvertOnTrainEnd
    - name: TestOnTrainEnd

  optimizer:
    name: SGD
    params:
      lr: 0.02

  scheduler:
    name: ConstantLR
```

The dataset in the example above is parsed using the LuxonisML parser. To learn more about it, refer to the [LuxonisParser
documentation](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-ml/luxonis-parser.md).

This example demonstrates the use of a predefined detection model, which you can explore further
[here](https://github.com/luxonis/luxonis-train/tree/main/luxonis_train/config/predefined_models#detectionmodel).

## 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](https://github.com/luxonis/luxonis-train/blob/main/luxonis_train/config/predefined_models/README.md).

## 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
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)
```

### Example: Custom Loss Function

```python
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()
```

### 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
model:
  nodes:
  - name: SegmentationHead
    losses:
    - name: CustomLoss
      params:
        smoothing: 0.0001

trainer:
  callbacks:
    - name: CustomCallback
      params:
        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

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

#### Using the Python API

```python
# Import your custom components first
from custom_components import *
from luxonis_train import LuxonisModel

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

Either method ensures your custom components are registered and available during the training process.
