Skip to content

Commit

Permalink
Tuning - Augmentation Subsets Support (#35)
Browse files Browse the repository at this point in the history 
Co-authored-by: Martin Kozlovsky <[email protected]>
  • Loading branch information
@klemen1999 @kozlov721
klemen1999 and kozlov721 authored Jun 19, 2024
1 parent bf69480 commit 88e8ff5
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 18 deletions.
34 changes: 18 additions & 16 deletions configs/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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>`. 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>`. 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:

Expand All @@ -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
Expand Down
9 changes: 9 additions & 0 deletions configs/example_tuning.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]
42 changes: 41 additions & 1 deletion luxonis_train/core/tuner.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os.path as osp
import random
from logging import getLogger
from typing import Any

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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
3 changes: 2 additions & 1 deletion luxonis_train/utils/config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
]


Expand Down

0 comments on commit 88e8ff5

Please sign in to comment.