概念

概述

在深入了解数据准备和模型训练之前，理解使 LuxonisTrain 高度灵活且可定制的核心概念至关重要。本节将带您了解配置和自定义的基础知识。

配置

LuxonisTrain 使用 YAML 配置文件来定义训练流程的每个方面，从模型架构到优化设置，为您配置训练过程提供了灵活性。本页将指导您了解各种可用的配置选项，确保您可以根据需要定制工具。有关所有参数、可能取值和默认设置的完整列表，请访问 LuxonisTrain 配置参考。

配置组件

模型（Model）：定义模型架构，包括节点、损失函数、可视化器和指标。
加载器（Loader）：通过选择合适的数据集来配置数据加载。有关如何创建数据集的更多详情，请参阅数据准备部分。
训练器（Trainer）：指定训练参数，包括批量大小、epoch 数、预处理（归一化、数据增强）、优化器、调度器、回调等。更多详情请参阅训练部分。
跟踪器（Tracker）：监控并记录训练指标以供分析。更多详情请参阅跟踪部分。
导出器（Exporter）：定义训练后模型的导出、存档和转换设置。更多详情请参阅导出部分。
调优器（Tuner）：调优超参数以优化模型性能。更多详情请参阅超参数调优部分。

示例配置文件

以下是一个示例 YAML 配置文件，帮助您了解如何构建 LuxonisTrain 设置：

Yaml

1model:
2  name: model_name
3
4  # 使用预定义的检测模型而不是手动定义架构
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# 数据集配置：从 Roboflow 下载并解析 COCO 数据集
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    # 使用 ImageNet 标准进行图像归一化
29    normalize:
30      active: true
31
32    # 使用 Albumentations 库进行图像增强
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

上例中的数据集是使用 LuxonisML 解析器解析的。要了解更多信息，请参阅 LuxonisParser 文档。此示例演示了如何使用预定义检测模型，您可以进一步探索此处。

预定义模型

您无需从头开始设计架构；可以根据要解决的任务使用预定义模型。这些模型经过充分测试，速度快且准确。您可以查看所有预定义模型此处。

自定义

LuxonisTrain 的模块化架构允许您自定义和扩展框架以满足特定需求。您可以自定义加载器（Loaders）、节点（Nodes）、损失函数（Losses）、指标（Metrics）、可视化器（Visualizers）、回调（Callbacks）、优化器（Optimizers）和调度器（Schedulers）等组件。此外，训练策略也可以完全自定义，允许您为不同的参数组设置不同的学习率，对不同的权重应用不同的调度器，并创建复杂的预热阶段。这种精细控制可实现复杂的优化策略，显著提高模型收敛性和性能。

自定义组件

要实现自定义组件，请继承相应的基类并注册它。以下是一些示例：

示例：自定义回调

Python

1import lightning.pytorch as pl
2
3from luxonis_train import LuxonisLightningModule
4from luxonis_train.registry import CALLBACKS
5
6@CALLBACKS.register()
7class CustomCallback(pl.Callback):
8    def __init__(self, message: str, **kwargs):
9        super().__init__(**kwargs)
10        self.message = message
11
12    # 将在每个训练 epoch 结束时被调用。
13    # 请查阅 PyTorch Lightning 文档以了解更多回调方法。
14    def on_train_epoch_end(
15        self,
16        trainer: pl.Trainer,
17        pl_module: LuxonisLightningModule,
18    ) -> None:
19        print(self.message)

示例：自定义损失函数

Python

1from torch import Tensor
2
3from luxonis_train import BaseLoss, Tasks
4
5# `BaseNode`、`BaseLoss`、`BaseMetric`
6# 和 `BaseVisualizer` 的子类会自动注册。
7class CustomLoss(BaseLoss):
8    supported_tasks = [Tasks.CLASSIFICATION, Tasks.SEGMENTATION]
9
10    def __init__(self, smoothing: float, **kwargs):
11        super().__init__(**kwargs)
12        self.smoothing = smoothing
13
14    def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
15        # 在此实现实际的损失逻辑。
16        value = predictions.sum() * self.smoothing
17        return value.abs()

在配置中使用自定义组件

要在训练配置中使用自定义组件，请在 YAML 文件中指定组件名称和参数。以下是一个示例：

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!"

在训练中包含自定义组件

创建自定义组件后，您需要在开始训练之前导入它们。有两种方法可以实现：

使用 CLI

Command Line

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

使用 Python API

Python

1# 首先导入自定义组件
2from custom_components import *
3from luxonis_train import LuxonisModel
4
5model = LuxonisModel("config.yaml")
6model.train()

两种方法都可确保您的自定义组件在训练过程中被注册并可用。

本页目录

概念

概述

配置

配置组件

示例配置文件

预定义模型

自定义

自定义组件

示例：自定义回调

示例：自定义损失函数

在配置中使用自定义组件

在训练中包含自定义组件

使用 CLI

使用 Python API