Merge branch 'main' into fix/resume-training-ema

Author: kozlov721

README.md

@@ -144,6 +144,10 @@ luxonis_train <command> --help
 
 Specific usage examples can be found in the respective sections below.
 
+> [!NOTE]
+> CLI commands `train`, `test`, and `tune` can be run with `--debug`
+> flag which allows the model to be used without a functional dataset.
+
 <a name="configuration"></a>
 
 ## ⚙️ Configuration

luxonis_train/__main__.py

@@ -6,21 +6,25 @@
 from functools import lru_cache
 from importlib.metadata import version
 from pathlib import Path
-from typing import TYPE_CHECKING, Annotated, Literal
+from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias
 
 import yaml
 from cyclopts import App, Group, Parameter, validators
 from loguru import logger
 from luxonis_ml.typing import Params, PathType
 
-from luxonis_train.config import Config
 from luxonis_train.upgrade import upgrade_config, upgrade_installation
 
+OptsType: TypeAlias = Annotated[
+    list[str] | None, Parameter(json_list=False, json_dict=False)
+]
+
 if TYPE_CHECKING:
     import numpy as np
 
     from luxonis_train import LuxonisModel
 
+
 app = App(
     help="Luxonis Train CLI",
     version=lambda: f"LuxonisTrain v{version('luxonis_train')}",
@@ -42,48 +46,23 @@ def create_model(
     config: PathType | Params | None,
     opts: list[str] | None = None,
     weights: PathType | None = None,
-    debug_mode: bool = False,
-    load_dataset_metadata: bool = True,
+    allow_empty_dataset: bool = False,
 ) -> "LuxonisModel":
     importlib.reload(sys.modules["luxonis_train"])
-    import torch
 
     from luxonis_train import LuxonisModel
-    from luxonis_train.utils.dataset_metadata import DatasetMetadata
-
-    if weights is not None and config is None:
-        ckpt = torch.load(weights, map_location="cpu")  # nosemgre
-        if "config" not in ckpt:  # pragma: no cover
-            raise ValueError(
-                f"Checkpoint '{weights}' does not contain the 'config' key. "
-                "Cannot restore `LuxonisModel` from checkpoint."
-            )
-        cfg = Config.get_config(upgrade_config(ckpt["config"]), opts)
-        dataset_metadata = None
-        if load_dataset_metadata:
-            if "dataset_metadata" not in ckpt:
-                logger.error("Checkpoint does not contain dataset metadata.")
-            else:
-                try:
-                    dataset_metadata = DatasetMetadata(
-                        **ckpt["dataset_metadata"]
-                    )
-                except Exception as e:  # pragma: no cover
-                    logger.error(
-                        "Failed to load dataset metadata from the checkpoint. "
-                        f"Error: {e}"
-                    )
-
-        return LuxonisModel(
-            cfg, debug_mode=debug_mode, dataset_metadata=dataset_metadata
-        )
 
-    return LuxonisModel(config, opts, debug_mode=debug_mode)
+    return LuxonisModel(
+        config,
+        opts,
+        weights=weights,
+        allow_empty_dataset=allow_empty_dataset,
+    )
 
 
 @app.command(group=training_group, sort_key=1)
 def train(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -99,29 +78,44 @@ def train(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
     @type debug: bool
-    @param debug: If True, the training will run in debug mode which
-        suppresses some exceptions to allow training without a fully
-        defined model.
+    @param debug: If true, allows the model to be constructed without
+    a valid dataset by setting `allow_empty_dataset` to True. This can
+    be useful for quick testing of the training loop.
     """
-    create_model(config, opts, weights, debug_mode=debug).train(
-        weights=weights
-    )
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=debug
+    ).train(weights=weights)
 
 
 @app.command(group=training_group, sort_key=2)
-def tune(opts: list[str] | None = None, /, *, config: str | None = None):
+def tune(
+    opts: OptsType = None,
+    /,
+    *,
+    config: str | None = None,
+    weights: str | None = None,
+    debug: bool = False,
+):
     """Start hyperparameter tuning.
 
     @type config: str
     @param config: Path to the configuration file.
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
+    @type weights: str
+    @param weights: Path to the model weights.
+    @type debug: bool
+    @param debug: If true, allows the model to be constructed without
+    a valid dataset by setting `allow_empty_dataset` to True. This can
+    be useful for quick testing of the tuning.
     """
-    create_model(config, opts).tune()
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=debug
+    ).tune()
 
 
 def _yield_visualizations(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     config: str | None = None,
     view: Literal["train", "val", "test"] = "train",
     size_multiplier: Annotated[
@@ -191,7 +185,7 @@ def get_visualization_item(
 
 @app.command(group=training_group, sort_key=3)
 def inspect(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -240,7 +234,7 @@ def get_window() -> str:
 
 @app.command(group=evaluation_group, sort_key=1)
 def test(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -261,18 +255,18 @@ def test(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
     @type debug: bool
-    @param debug: If True, the training will run in debug mode which
-        suppresses some exceptions to allow training without a fully
-        defined model.
+    @param debug: If true, allows the model to be constructed without
+    a valid dataset by setting `allow_empty_dataset` to True. This can
+    be useful for quick testing of the evaluation loop.
     """
-    create_model(config, opts, weights, debug_mode=debug).test(
-        view=view, weights=weights
-    )
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=debug
+    ).test(view=view, weights=weights)
 
 
 @app.command(group=evaluation_group, sort_key=2)
 def infer(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -302,7 +296,9 @@ def infer(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
     """
-    create_model(config, opts, weights=weights, debug_mode=True).infer(
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=True
+    ).infer(
         view=view,
         save_dir=save_dir,
         source_path=source_path,
@@ -312,7 +308,7 @@ def infer(
 
 @app.command(group=annotation_group, sort_key=0)
 def annotate(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     dir_path: Path,
@@ -323,7 +319,6 @@ def annotate(
     delete_local: bool = True,
     delete_remote: bool = True,
     team_id: str | None = None,
-    debug: bool = False,
 ):
     """Run annotation on a custom directory of images.
 
@@ -353,11 +348,7 @@ def annotate(
     @param opts: A list of optional CLI overrides of the config file.
     """
     model = create_model(
-        config,
-        opts,
-        weights=weights,
-        load_dataset_metadata=True,
-        debug_mode=debug,
+        config, opts, weights=weights, allow_empty_dataset=True
     )
 
     model.annotate(
@@ -373,7 +364,7 @@ def annotate(
 
 @app.command(group=export_group, sort_key=1)
 def export(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -400,14 +391,14 @@ def export(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the
     """
-    create_model(config, opts, weights=weights, debug_mode=True).export(
-        save_path=save_path, weights=weights, ckpt_only=ckpt_only
-    )
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=True
+    ).export(save_path=save_path, weights=weights, ckpt_only=ckpt_only)
 
 
 @app.command(group=export_group, sort_key=2)
 def archive(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None,
@@ -426,14 +417,14 @@ def archive(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
     """
-    create_model(str(config), opts, weights=weights).archive(
-        path=executable, weights=weights
-    )
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=True
+    ).archive(path=executable, weights=weights)
 
 
 @app.command(group=export_group, sort_key=3)
 def convert(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     config: str | None = None,
@@ -456,9 +447,9 @@ def convert(
     @type opts: list[str]
     @param opts: A list of optional CLI overrides of the config file.
     """
-    create_model(config, opts, weights=weights).convert(
-        weights=weights, save_dir=save_dir
-    )
+    create_model(
+        config, opts, weights=weights, allow_empty_dataset=True
+    ).convert(save_dir=save_dir, weights=weights)
 
 
 @upgrade_app.command()
@@ -500,7 +491,7 @@ def config(
 
 @upgrade_app.command(name=["checkpoint", "ckpt"])
 def checkpoint(
-    opts: list[str] | None = None,
+    opts: OptsType = None,
     /,
     *,
     path: Annotated[
@@ -518,11 +509,10 @@ def checkpoint(
     @param new: Where to save the upgraded checkpoint. If left empty,
         the old file will be overriden.
     """
+    from luxonis_train import LuxonisModel
+
     logger.info("Performing a full checkpoint upgrade.")
-    cfg = None
-    if config is not None:
-        cfg = upgrade_config(config)
-    model = create_model(config=cfg, weights=path, opts=opts, debug_mode=True)
+    model = LuxonisModel(config, opts, weights=path, allow_empty_dataset=True)
     model.lightning_module.load_checkpoint(path)
 
     # Needs to be called in order to attach the model to the trainer

luxonis_train/callbacks/ema.py

@@ -298,7 +298,7 @@ def on_save_checkpoint(
         trainer: pl.Trainer,
         pl_module: pl.LightningModule,
         checkpoint: dict,
-    ) -> dict[str, Any] | None:
+    ) -> None:
         """Save the EMA state dictionary into the checkpoint.
 
         @type trainer: L{pl.Trainer}

luxonis_train/callbacks/upload_checkpoint.py

@@ -1,10 +1,12 @@
+from copy import copy
 from pathlib import Path
 from typing import Any
 
 import lightning.pytorch as pl
 import torch
 from lightning.pytorch.callbacks import ModelCheckpoint
 from loguru import logger
+from typing_extensions import override
 
 import luxonis_train as lxt
 from luxonis_train.registry import CALLBACKS
@@ -24,12 +26,15 @@ def __init__(self):
         super().__init__()
         self.last_best_checkpoints = set()
 
+    @override
     def on_save_checkpoint(
         self,
         trainer: pl.Trainer,
         module: "lxt.LuxonisLightningModule",
         checkpoint: dict[str, Any],
     ) -> None:
+        checkpoint = copy(checkpoint)
+        module._add_custom_data_to_checkpoint(checkpoint)
         checkpoint_paths = [
             c.best_model_path
             for c in trainer.checkpoint_callbacks