Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Tuning - Augmentation Subsets Support #35

Merged
merged 8 commits into from
Jun 19, 2024
Merged

Tuning - Augmentation Subsets Support #35

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e3e915e
added tuner subset sampling for augmentations
klemen1999 May 31, 2024
0269143
updated README
klemen1999 May 31, 2024
d38d395
formatting
klemen1999 May 31, 2024
2a40f14
updated readme
klemen1999 Jun 2, 2024
a6ad01b
cleaner aug name index mapping
klemen1999 Jun 7, 2024
3242631
Merge branch 'dev' into feature/aug_subset_tune
kozlov721 Jun 13, 2024
85f267e
fixed CLI archiver type
kozlov721 Jun 13, 2024
4abe9f8
Merge branch 'dev' into feature/aug_subset_tune
kozlov721 Jun 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 18 additions & 16 deletions configs/README.md
View file
Open in desktop
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
Open in desktop
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]
2 changes: 1 addition & 1 deletion luxonis_train/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,7 @@ def inspect(
@app.command()
def archive(
executable: Annotated[
Optional[Path], typer.Option(help="Path to the model file.", show_default=False)
str, typer.Option(help="Path to the model file.", show_default=False)
],
config: ConfigType = None,
opts: OptsType = None,
Expand Down
42 changes: 41 additions & 1 deletion luxonis_train/core/tuner.py
View file
Open in desktop
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 @@ -121,8 +122,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 @@ -189,6 +197,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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How are we preventing it from selecting a subset that was already selected in a previous run?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For this question I assume that typically we would use the augmentations tuning with all the other tunings disabled to get purely augmentation results. Would that be correct?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, there isn't really any prevention in place for selecting same subset in subsequent runs. We could change it to first get all possible combinations and then loop through them each run. But then number of trials should also be set to number of combinations so we go through all of them. I guess it depends on what would be the typical usecase for this aug subset sampling.

As you mentioned below, perhaps looping over whole powerset is a more sensible impelmentation since normally we shouldn't be limited by the number of augmentations used and instead use all of those that produce better results. But to see if augmentation works you normally have to train the model for more epochs (at the beginning hard augs could produce worse results but with time perhaps those are the ones that actually improve acc on test set) and going over whole powerset could take a while (maybe add ability to limit minimum size of the subset to prune smaller ones).
CC: @tersekmatija on what would be the intended usecase of aug subsampling.

Copy link
Collaborator

@kozlov721 kozlov721 Jun 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When I'm thinking more about the augmentation tuning, I think this will actually be more complex and will need more optimization. For now, I can think of a few optimizations if we'll work with some assumptions about the augmentations' effect on the training.

My idea: I think we can assume the order of augmentations should not (reasonably) matter, so if augmentation $A$ increases the model's performance, it doesn't matter where in the pipeline it's placed so we can lock it in place and continue only with the remaining augmentations and one-smaller subset size.
This would decrease the number of possibilities from $n^k$ to $n \choose k$ ($n$ being the number of all augmentations and $k$ being the subset size), but that's still way too many for any reasonable usage. We support ~80 augmentations and even with a subset of size only 5, there's 22 milion combinations to go through.

However, as an extension of the above, if augmentation $A$ did not increase the performance, we can discard it for the rest of the tuning. That's because if $A$ didn't improve, a different augmentation $B$ did improve, and the order doesn't matter, then $A \circ B \simeq B \simeq B \circ A$, so there's no need to try $A$ ever again.
This would bring the number of combinations all the way down to $n$, even with tunable subset size.

More digestible in code:

best_augmentations = []
for a in all_augmentations:
    if improves(a):
        best_augmentations.append(a)
return best_augmentations

Someone would need to double-check my math on this though (and test whether the assumption even holds in the first place).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Math looks correct :) The challenge here is if you use more augmentations at once. If you use a and b and they improve the accuracy, you don't know whether contribution was made by a or b. This means that in any case you need to run at least n+1 (where n=80) combinations to find out which show improvement and which do not. Furthermore, I think the challenge then might be that single augmentations improve the performance, but a combination of them decreases it (image gets too corrupted to be useful).

I still think it should be up to the user to define which augmentations are reasonable, and then merely test a few combinations to check whether they are useful or not. I think in the above scenario it's also hard to define what k is?

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 @@ -221,3 +241,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
Open in desktop
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
Loading