diff --git a/configs/README.md b/configs/README.md index e6eb0bad..6901375c 100644 --- a/configs/README.md +++ b/configs/README.md @@ -166,14 +166,14 @@ Here you can change everything related to actual training of the model. We use [Albumentations](https://albumentations.ai/docs/) library for `augmentations`. [Here](https://albumentations.ai/docs/api_reference/full_reference/#pixel-level-transforms) you can see a list of all pixel level augmentations supported, and [here](https://albumentations.ai/docs/api_reference/full_reference/#spatial-level-transforms) you see all spatial level transformations. In config you can specify any augmentation from this lists and their params. Additionaly we support `Mosaic4` batch augmentation and letterbox resizing if `keep_aspect_ratio: True`. -| Key | Type | Default value | Description | -| ----------------- | ------------------------------------------------------------------------------------ | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| train_image_size | list\[int\] | \[256, 256\] | image size used for training \[height, width\] | -| keep_aspect_ratio | bool | True | bool if keep aspect ration while resizing | -| train_rgb | bool | True | bool if train on rgb or bgr | -| normalize.active | bool | True | bool if use normalization | -| normalize.params | dict | {} | params for normalization, see [documentation](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | -| augmentations | list\[{"name": Name of the augmentation, "params": Parameters of the augmentation}\] | \[\] | list of Albumentations augmentations | +| Key | Type | Default value | Description | +| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| train_image_size | list\[int\] | \[256, 256\] | image size used for training \[height, width\] | +| keep_aspect_ratio | bool | True | bool if keep aspect ration while resizing | +| train_rgb | bool | True | bool if train on rgb or bgr | +| normalize.active | bool | True | bool if use normalization | +| normalize.params | dict | {} | params for normalization, see [documentation](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize) | +| augmentations | list\[{"name": Name of the augmentation, "active": Bool if aug is active, by default set to True, "params": Parameters of the augmentation}\] | \[\] | list of Albumentations augmentations | ### Optimizer @@ -241,14 +241,15 @@ Option specific for ONNX export. Here you can specify options for tuning. -| Key | Type | Default value | Description | -| ----------------------- | ----------------- | ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| study_name | str | "test-study" | Name of the study. | -| continue_existing_study | bool | True | Weather to continue existing study if `study_name` already exists. | -| use_pruner | bool | True | Whether to use the MedianPruner. | -| n_trials | int \| None | 15 | Number of trials for each process. `None` represents no limit in terms of numbner of trials. | -| timeout | int \| None | None | Stop study after the given number of seconds. | -| params | dict\[str, list\] | {} | Which parameters to tune. The keys should be in the format `key1.key2.key3_`. Type can be one of `[categorical, float, int, longuniform, uniform]`. For more information about the types, visit [Optuna documentation](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.trial.Trial.html). | +| Key | Type | Default value | Description | +| ---------- | ----------------- | ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| study_name | str | "test-study" | Name of the study. | +| use_pruner | bool | True | Whether to use the MedianPruner. | +| n_trials | int \| None | 15 | Number of trials for each process. `None` represents no limit in terms of numbner of trials. | +| timeout | int \| None | None | Stop study after the given number of seconds. | +| params | dict\[str, list\] | {} | Which parameters to tune. The keys should be in the format `key1.key2.key3_`. Type can be one of `[categorical, float, int, longuniform, uniform, subset]`. For more information about the types, visit [Optuna documentation](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.trial.Trial.html). | + +**Note**: "subset" sampling is currently only supported for augmentations. You can specify a set of augmentations defined in `trainer` to choose from and every run subset of random N augmentations will be active (`is_active` parameter will be True for chosen ones and False for the rest in the set). Example of params for tuner block: @@ -258,6 +259,7 @@ tuner: trainer.optimizer.name_categorical: ["Adam", "SGD"] trainer.optimizer.params.lr_float: [0.0001, 0.001] trainer.batch_size_int: [4, 16, 4] + trainer.preprocessing.augmentations_subset: [["Defocus", "Sharpen", "Flip"], 2] ``` ### Storage diff --git a/configs/example_tuning.yaml b/configs/example_tuning.yaml index 009abc41..9a8bfd79 100755 --- a/configs/example_tuning.yaml +++ b/configs/example_tuning.yaml @@ -21,6 +21,14 @@ trainer: keep_aspect_ratio: False normalize: active: True + augmentations: + - name: Defocus + params: + p: 0.1 + - name: Sharpen + params: + p: 0.1 + - name: Flip batch_size: 4 epochs: &epochs 10 @@ -38,3 +46,4 @@ tuner: trainer.optimizer.name_categorical: ["Adam", "SGD"] trainer.optimizer.params.lr_float: [0.0001, 0.001] trainer.batch_size_int: [4, 16, 4] + trainer.preprocessing.augmentations_subset: [["Defocus", "Sharpen", "Flip"], 2] diff --git a/luxonis_train/core/tuner.py b/luxonis_train/core/tuner.py index 4dfc780b..13d56ca4 100644 --- a/luxonis_train/core/tuner.py +++ b/luxonis_train/core/tuner.py @@ -1,4 +1,5 @@ import os.path as osp +import random from logging import getLogger from typing import Any @@ -125,8 +126,15 @@ def _objective(self, trial: optuna.trial.Trial) -> float: curr_params = self._get_trial_params(trial) curr_params["model.predefined_model"] = None + + cfg_copy = self.cfg.model_copy(deep=True) + cfg_copy.trainer.preprocessing.augmentations = [ + a + for a in cfg_copy.trainer.preprocessing.augmentations + if a.name != "Normalize" + ] # manually remove Normalize so it doesn't duplicate it when creating new cfg instance Config.clear_instance() - cfg = Config.get_config(self.cfg.model_dump(), curr_params) + cfg = Config.get_config(cfg_copy.model_dump(), curr_params) child_tracker.log_hyperparams(curr_params) @@ -193,6 +201,18 @@ def _get_trial_params(self, trial: optuna.trial.Trial) -> dict[str, Any]: key_name = "_".join(key_info[:-1]) key_type = key_info[-1] match key_type, value: + case "subset", [list(whole_set), int(subset_size)]: + if key_name.split(".")[-1] != "augmentations": + raise ValueError( + "Subset sampling currently only supported for augmentations" + ) + whole_set_indices = self._augs_to_indices(whole_set) + subset = random.sample(whole_set_indices, subset_size) + for aug_id in whole_set_indices: + new_params[f"{key_name}.{aug_id}.active"] = ( + True if aug_id in subset else False + ) + continue case "categorical", list(lst): new_value = trial.suggest_categorical(key_name, lst) case "float", [float(low), float(high), *tail]: @@ -225,3 +245,23 @@ def _get_trial_params(self, trial: optuna.trial.Trial) -> dict[str, Any]: "No paramteres to tune. Specify them under `tuner.params`." ) return new_params + + def _augs_to_indices(self, aug_names: list[str]) -> list[int]: + """Maps augmentation names to indices.""" + all_augs = [a.name for a in self.cfg.trainer.preprocessing.augmentations] + aug_indices = [] + for aug_name in aug_names: + if aug_name == "Normalize": + logger.warn( + f"'{aug_name}' should be tuned directly by adding '...normalize.active_categorical' to the tuner params, skipping." + ) + continue + try: + index = all_augs.index(aug_name) + aug_indices.append(index) + except ValueError: + logger.warn( + f"Augmentation '{aug_name}' not found under trainer augemntations, skipping." + ) + continue + return aug_indices diff --git a/luxonis_train/utils/config.py b/luxonis_train/utils/config.py index afd1f5c7..768c6f04 100644 --- a/luxonis_train/utils/config.py +++ b/luxonis_train/utils/config.py @@ -292,7 +292,8 @@ class TunerConfig(CustomBaseModel): timeout: int | None = None storage: StorageConfig = StorageConfig() params: Annotated[ - dict[str, list[str | int | float | bool]], Field(default={}, min_length=1) + dict[str, list[str | int | float | bool | list]], + Field(default={}, min_length=1), ]