Custom optimizer

README.md

@@ -567,6 +567,7 @@ model.tune()
 - [**Callbacks**](https://github.com/luxonis/luxonis-train/blob/main/luxonis_train/callbacks/README.md): Allow custom code to be executed at different stages of training.
 - [**Optimizers**](https://github.com/luxonis/luxonis-train/blob/main/configs/README.md#optimizer): Control how the model's weights are updated.
 - [**Schedulers**](https://github.com/luxonis/luxonis-train/blob/main/configs/README.md#scheduler): Adjust the learning rate during training.
+- [**Training Strategy**](https://github.com/luxonis/luxonis-train/blob/main/configs/README.md#training-strategy): Specify a custom combination of optimizer and scheduler to tailor the training process for specific use cases.
 
 **Creating Custom Components:**
 
@@ -581,6 +582,7 @@ Registered components can be referenced in the config file. Custom components ne
 - **Callbacks** - [`lightning.pytorch.callbacks.Callback`](https://lightning.ai/docs/pytorch/stable/extensions/callbacks.html), requires manual registration to the `CALLBACKS` registry
 - **Optimizers** - [`torch.optim.Optimizer`](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer), requires manual registration to the `OPTIMIZERS` registry
 - **Schedulers** - [`torch.optim.lr_scheduler.LRScheduler`](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate), requires manual registration to the `SCHEDULERS` registry
+- **Training Strategy** - [`BaseTrainingStrategy`](https://github.com/luxonis/luxonis-train/blob/main/luxonis_train/strategies/base_strategy.py)
 
 **Examples:**
 

configs/README.md

@@ -376,6 +376,37 @@ trainer:
       eta_min: 0
 ```
 
+### Training Strategy
+
+Defines the training strategy to be used. Currently, only the `TripleLRSGDStrategy` is supported, but more strategies will be added in the future.
+
+| Key               | Type    | Default value           | Description                                    |
+| ----------------- | ------- | ----------------------- | ---------------------------------------------- |
+| `name`            | `str`   | `"TripleLRSGDStrategy"` | Name of the training strategy                  |
+| `warmup_epochs`   | `int`   | `3`                     | Number of epochs for the warmup phase          |
+| `warmup_bias_lr`  | `float` | `0.1`                   | Learning rate for bias during the warmup phase |
+| `warmup_momentum` | `float` | `0.8`                   | Momentum value during the warmup phase         |
+| `lr`              | `float` | `0.02`                  | Initial learning rate                          |
+| `lre`             | `float` | `0.0002`                | End learning rate                              |
+| `momentum`        | `float` | `0.937`                 | Momentum for the optimizer                     |
+| `weight_decay`    | `float` | `0.0005`                | Weight decay value                             |
+| `nesterov`        | `bool`  | `true`                  | Use Nesterov momentum or not                   |
+
+**Example:**
+
+```yaml
+training_strategy:
+  name: "TripleLRSGDStrategy"
+  warmup_epochs: 3
+  warmup_bias_lr: 0.1
+  warmup_momentum: 0.8
+  lr: 0.02
+  lre: 0.0002
+  momentum: 0.937
+  weight_decay: 0.0005
+  nesterov: true
+```
+
 ## Exporter
 
 Here you can define configuration for exporting.

luxonis_train/__init__.py

@@ -10,6 +10,7 @@
     from .nodes import *
     from .optimizers import *
     from .schedulers import *
+    from .strategies import *
     from .utils import *
 except ImportError as e:
     warnings.warn(

luxonis_train/callbacks/__init__.py

@@ -25,6 +25,7 @@
 from .metadata_logger import MetadataLogger
 from .module_freezer import ModuleFreezer
 from .test_on_train_end import TestOnTrainEnd
+from .training_manager import TrainingManager
 from .upload_checkpoint import UploadCheckpoint
 
 CALLBACKS.register_module(module=EarlyStopping)
@@ -38,6 +39,7 @@
 CALLBACKS.register_module(module=ModelPruning)
 CALLBACKS.register_module(module=GradCamCallback)
 CALLBACKS.register_module(module=EMACallback)
+CALLBACKS.register_module(module=TrainingManager)
 
 
 __all__ = [
@@ -53,4 +55,5 @@
     "GPUStatsMonitor",
     "GradCamCallback",
     "EMACallback",
+    "TrainingManager",
 ]

luxonis_train/callbacks/training_manager.py

@@ -0,0 +1,28 @@
+import pytorch_lightning as pl
+
+from luxonis_train.strategies.base_strategy import BaseTrainingStrategy
+
+
+class TrainingManager(pl.Callback):
+    def __init__(self, strategy: BaseTrainingStrategy | None = None):
+        """Training manager callback that updates the parameters of the
+        training strategy.
+
+        @type strategy: BaseTrainingStrategy
+        @param strategy: The strategy to be used.
+        """
+        self.strategy = strategy
+
+    def on_after_backward(
+        self, trainer: pl.Trainer, pl_module: pl.LightningModule
+    ):
+        """PyTorch Lightning hook that is called after the backward
+        pass.
+
+        @type trainer: pl.Trainer
+        @param trainer: The trainer object.
+        @type pl_module: pl.LightningModule
+        @param pl_module: The pl_module object.
+        """
+        if self.strategy is not None:
+            self.strategy.update_parameters(pl_module)

luxonis_train/config/config.py

@@ -336,6 +336,11 @@ class SchedulerConfig(BaseModelExtraForbid):
     params: Params = {}
 
 
+class TrainingStrategyConfig(BaseModelExtraForbid):
+    name: str
+    params: Params = {}
+
+
 class TrainerConfig(BaseModelExtraForbid):
     preprocessing: PreprocessingConfig = PreprocessingConfig()
     use_rich_progress_bar: bool = True
@@ -382,6 +387,7 @@ class TrainerConfig(BaseModelExtraForbid):
 
     optimizer: OptimizerConfig = OptimizerConfig()
     scheduler: SchedulerConfig = SchedulerConfig()
+    training_strategy: TrainingStrategyConfig | None = None
 
     @model_validator(mode="after")
     def validate_deterministic(self) -> Self:

luxonis_train/models/luxonis_lightning.py

@@ -25,7 +25,11 @@
     combine_visualizations,
     get_denormalized_images,
 )
-from luxonis_train.callbacks import BaseLuxonisProgressBar, ModuleFreezer
+from luxonis_train.callbacks import (
+    BaseLuxonisProgressBar,
+    ModuleFreezer,
+    TrainingManager,
+)
 from luxonis_train.config import AttachedModuleConfig, Config
 from luxonis_train.nodes import BaseNode
 from luxonis_train.utils import (
@@ -42,6 +46,7 @@
     CALLBACKS,
     OPTIMIZERS,
     SCHEDULERS,
+    STRATEGIES,
     Registry,
 )
 
@@ -268,6 +273,24 @@
 
         self.load_checkpoint(self.cfg.model.weights)
 
+        if self.cfg.trainer.training_strategy is not None:
+            if self.cfg.trainer.optimizer is not None:
+                logger.warning(
+                    "Training strategy is active; the specified optimizer will be ignored."
+                )
+            if self.cfg.trainer.scheduler is not None:
+                logger.warning(
+                    "Training strategy is active; the specified scheduler will be ignored."
+                )
+            self.training_strategy = STRATEGIES.get(
+                self.cfg.trainer.training_strategy.name
+            )(
+                pl_module=self,
+                params=self.cfg.trainer.training_strategy.params,  # type: ignore
+            )
+        else:
+            self.training_strategy = None
+
     @property
     def core(self) -> "luxonis_train.core.LuxonisModel":
         """Returns the core model."""
@@ -849,6 +872,9 @@
                     CALLBACKS.get(callback.name)(**callback.params)
                 )
 
+        if self.training_strategy is not None:
+            callbacks.append(TrainingManager(strategy=self.training_strategy))  # type: ignore
+
         return callbacks
 
     def configure_optimizers(
@@ -858,6 +884,9 @@
         list[torch.optim.lr_scheduler.LRScheduler],
     ]:
         """Configures model optimizers and schedulers."""
+        if self.training_strategy is not None:
+            return self.training_strategy.configure_optimizers()
+
         cfg_optimizer = self.cfg.trainer.optimizer
         cfg_scheduler = self.cfg.trainer.scheduler
 

luxonis_train/nodes/backbones/efficientrep/efficientrep.py

@@ -125,9 +125,23 @@ def __init__(
             )
         )
 
+        self.initialize_weights()
+
         if download_weights and var.weights_path:
             self.load_checkpoint(var.weights_path)
 
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                pass
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eps = 0.001
+                m.momentum = 0.03
+            elif isinstance(
+                m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU)
+            ):
+                m.inplace = True
+
     def set_export_mode(self, mode: bool = True) -> None:
         """Reparametrizes instances of L{RepVGGBlock} in the network.
 

luxonis_train/nodes/blocks/blocks.py

@@ -56,6 +56,19 @@ def __init__(self, n_classes: int, in_channels: int):
 
         prior_prob = 1e-2
         self._initialize_weights_and_biases(prior_prob)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                pass
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eps = 0.001
+                m.momentum = 0.03
+            elif isinstance(
+                m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU)
+            ):
+                m.inplace = True
 
     def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
         out_feature = self.decoder(x)

luxonis_train/nodes/heads/efficient_bbox_head.py

@@ -95,12 +95,26 @@ def __init__(
                 f"output{i+1}_yolov6r2" for i in range(self.n_heads)
             ]
 
+        self.initialize_weights()
+
         if download_weights:
             # TODO: Handle variants of head in a nicer way
             if self.in_channels == [32, 64, 128]:
                 weights_path = "https://github.com/luxonis/luxonis-train/releases/download/v0.1.0-beta/efficientbbox_head_n_coco.ckpt"
                 self.load_checkpoint(weights_path, strict=False)
 
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                pass
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eps = 0.001
+                m.momentum = 0.03
+            elif isinstance(
+                m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU)
+            ):
+                m.inplace = True
+
     def forward(
         self, inputs: list[Tensor]
     ) -> tuple[list[Tensor], list[Tensor], list[Tensor]]:

luxonis_train/nodes/necks/reppan_neck/reppan_neck.py

@@ -165,9 +165,23 @@ def __init__(
             out_channels = channels_list_down_blocks[2 * i + 1]
             curr_n_repeats = n_repeats_down_blocks[i]
 
+        self.initialize_weights()
+
         if download_weights and var.weights_path:
             self.load_checkpoint(var.weights_path)
 
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                pass
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eps = 0.001
+                m.momentum = 0.03
+            elif isinstance(
+                m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU)
+            ):
+                m.inplace = True
+
     def forward(self, inputs: list[Tensor]) -> list[Tensor]:
         x = inputs[-1]
         up_block_outs: list[Tensor] = []

luxonis_train/strategies/__init__.py

@@ -0,0 +1,7 @@
+from .base_strategy import BaseTrainingStrategy
+from .triple_lr_sgd import TripleLRScheduler
+
+__all__ = [
+    "TripleLRScheduler",
+    "BaseTrainingStrategy",
+]

luxonis_train/strategies/base_strategy.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+
+import pytorch_lightning as pl
+from luxonis_ml.utils.registry import AutoRegisterMeta
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+from luxonis_train.utils.registry import STRATEGIES
+
+
+class BaseTrainingStrategy(
+    ABC,
+    metaclass=AutoRegisterMeta,
+    register=False,
+    registry=STRATEGIES,
+):
+    def __init__(self, pl_module: pl.LightningModule):
+        self.pl_module = pl_module
+
+    @abstractmethod
+    def configure_optimizers(
+        self,
+    ) -> tuple[list[Optimizer], list[LRScheduler]]:
+        pass
+
+    @abstractmethod
+    def update_parameters(self, *args, **kwargs):
+        pass