Optuna HPO

`run_optuna(cfg)` ¶

Run an Optuna study from a configured TrainerConfig.

This is the recommended user-facing entrypoint. It creates an OptunaRunner(cfg), calls runner.run(), and returns the resulting Optuna Study object.

Parameters:

Name	Type	Description	Default
`cfg`	`TrainerConfig`	Base training config. `cfg.optuna` must be set to an `OptunaConfig`.	required

Returns:

Name	Type	Description
`Any`	`typing.Any`	The Optuna `Study` object returned by `OptunaRunner.run()`.

Usage

from deepvisiontools.train import TrainerConfig
from deepvisiontools.train.optuna_wrapper import (
    OptunaConfig,
    SearchSpaceSpec,
    run_optuna,
)

cfg = TrainerConfig.build_registered_config(
    model="detector",
    config="detector_convnextv2-base_vfnet",
    train_loader=trainloader,
    valid_loader=validloader,
)

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="vfnet_hpo",
    storage="vfnet_hpo.db",
    n_trials=25,
    search_space=[
        SearchSpaceSpec(
            name="lr",
            kind="float",
            path="optim.optim_args.lr",
            low=1e-5,
            high=1e-3,
            log=True,
        ),
    ],
)

study = run_optuna(cfg)

print("Best value:", study.best_value)
print("Best params:", study.best_params)

Source code in src/deepvisiontools/train/optuna_wrapper.py

def run_optuna(cfg: TrainerConfig) -> Any:
    """
    Run an Optuna study from a configured `TrainerConfig`.

    This is the recommended user-facing entrypoint. It creates an
    `OptunaRunner(cfg)`, calls `runner.run()`, and returns the resulting Optuna
    `Study` object.

    Args:
        cfg (`TrainerConfig`): Base training config. `cfg.optuna` must be set to
            an `OptunaConfig`.

    Returns:
        Any: The Optuna `Study` object returned by `OptunaRunner.run()`.

    ???+ example "Usage"
        ```python
        from deepvisiontools.train import TrainerConfig
        from deepvisiontools.train.optuna_wrapper import (
            OptunaConfig,
            SearchSpaceSpec,
            run_optuna,
        )

        cfg = TrainerConfig.build_registered_config(
            model="detector",
            config="detector_convnextv2-base_vfnet",
            train_loader=trainloader,
            valid_loader=validloader,
        )

        cfg.optuna = OptunaConfig(
            enabled=True,
            study_name="vfnet_hpo",
            storage="vfnet_hpo.db",
            n_trials=25,
            search_space=[
                SearchSpaceSpec(
                    name="lr",
                    kind="float",
                    path="optim.optim_args.lr",
                    low=1e-5,
                    high=1e-3,
                    log=True,
                ),
            ],
        )

        study = run_optuna(cfg)

        print("Best value:", study.best_value)
        print("Best params:", study.best_params)
        ```
    """
    return OptunaRunner(cfg).run()

`OptunaRunner(base_cfg)` ¶

Study-level Optuna orchestration wrapper.

OptunaRunner owns the Optuna study lifecycle. It does not replace Trainer: it clones a base TrainerConfig for each trial, applies search mutations, runs one training job, and returns one objective value to Optuna.

Parameters:

Name	Type	Description	Default
`base_cfg`	`TrainerConfig`	Base training configuration. It must have `base_cfg.optuna` set to an `OptunaConfig`.	required

What happens for each trial

Clone the base TrainerConfig.
Force trial logging into optuna.mlflow_experiment_name.
Optionally disable checkpoints and model exports for HPO.
Apply search_space suggestions to dotted config paths.
Apply built-in augmentation_space mutations.
Apply optional trial_mutator_target and data_mutator_target.
Run either an in-process training job or a distributed child job.
Record monitored values after epochs.
Reduce the history into a final objective with OptunaScoreConfig.

Usage

from deepvisiontools.train.optuna_wrapper import OptunaConfig, OptunaRunner

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="my_hpo",
    storage="my_hpo.db",
    n_trials=30,
)

runner = OptunaRunner(cfg)
study = runner.run()

print(study.best_value)
print(study.best_params)

Custom trial and data mutators

# my_project/hpo_callbacks.py

def mutate_trial(trial, cfg):
    # Use this escape hatch when a simple dotted path is not enough.
    if trial.suggest_categorical("small_batch", [False, True]):
        cfg.precision.accumulate_steps = 4
        cfg.devices.num_workers = 2

def rebuild_data(trial, cfg):
    # Rebuild loaders/datasets here if the search changes dataset-level
    # objects that cannot be expressed by AugmentationPipelineSpec.
    batch_size = trial.suggest_categorical("batch_size", [2, 4, 8])
    cfg.train_loader.batch_size = batch_size
    cfg.valid_loader.batch_size = batch_size

from deepvisiontools.train.optuna_wrapper import OptunaConfig, run_optuna

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="custom_mutator_hpo",
    storage="custom_mutator_hpo.db",
    n_trials=20,
    trial_mutator_target="my_project.hpo_callbacks:mutate_trial",
    data_mutator_target="my_project.hpo_callbacks:rebuild_data",
)

study = run_optuna(cfg)

Source code in src/deepvisiontools/train/optuna_wrapper.py

def __init__(self, base_cfg: TrainerConfig):
    self.base_cfg = base_cfg
    optuna_cfg = getattr(base_cfg, "optuna", None)
    if optuna_cfg is None:
        raise ValueError("OptunaRunner requires base_cfg.optuna to be configured.")
    self.optuna_cfg: OptunaConfig = optuna_cfg
    if self.optuna_cfg is None:
        raise ValueError(
            "base_cfg.optuna must be set before creating OptunaRunner."
        )
    self.optuna_cfg.validate()

OptunaConfig(enabled=False, study_name='dvt_optuna_study', storage='optuna_study.db', load_if_exists=True, n_trials=None, timeout_seconds=None, n_jobs=1, gc_after_trial=True, mlflow_experiment_name='optuna_run', mlflow_run_name_template='trial_{trial_number:05d}', register_study_in_mlflow=True, disable_checkpointing=True, distributed_trials=False, trial_devices='auto', parent_poll_interval_seconds=0.5, decision_wait_seconds=20.0, child_exit_grace_seconds=60.0, keep_trial_workdirs=False, sampler=OptunaSamplerConfig(), pruner=OptunaPrunerConfig(), score=OptunaScoreConfig(), plateau_stop=PlateauStopConfig(), search_space=list(), augmentation_space=list(), trial_mutator_target=None, data_mutator_target=None) `dataclass` ¶

Top-level Optuna configuration used by OptunaRunner and run_optuna.

Trainer remains a single-training-run object. Optuna is handled one level above it: for each trial, OptunaRunner clones the base TrainerConfig, applies search-space suggestions and optional mutators, runs training, then reduces the monitored history into one objective value.

Parameters:

Name	Type	Description	Default
`enabled`	`bool`	User-facing switch/metadata flag for YAML/config files. The actual entrypoint is `run_optuna(cfg)`. Defaults to False.	`False`
`study_name`	`str`	Optuna study name. Defaults to `"dvt_optuna_study"`.	`'dvt_optuna_study'`
`storage`	`str \| None`	Optuna storage. None creates an in-memory study. A bare path such as `"optuna_study.db"` is converted to a local SQLite URL. Full URLs such as `"sqlite:///path/to/study.db"` are also accepted.	`'optuna_study.db'`
`load_if_exists`	`bool`	Reuse an existing study with the same name and storage. Defaults to True.	`True`
`n_trials`	`int \| None`	Maximum number of trials. If None, Optuna uses only other stop criteria such as timeout or callbacks.	`None`
`timeout_seconds`	`int \| None`	Maximum optimization time in seconds.	`None`
`n_jobs`	`int`	Number of Optuna trials run in parallel by the parent process. Defaults to 1. For GPU training, 1 is usually safest unless you intentionally manage resources.	`1`
`gc_after_trial`	`bool`	Forwarded to `study.optimize`; also complements deepvisiontools cleanup. Defaults to True.	`True`
`mlflow_experiment_name`	`str`	MLflow experiment used for trial runs and optional study registration. Defaults to `"optuna_run"`.	`'optuna_run'`
`mlflow_run_name_template`	`str`	Trial run-name template. It can use `{trial_number}`. Defaults to `"trial_{trial_number:05d}"`.	`'trial_{trial_number:05d}'`
`register_study_in_mlflow`	`bool`	If True, logs study metadata in MLflow. Defaults to True.	`True`
`disable_checkpointing`	`bool`	If True, disables training checkpoints and portable model artifact export inside trials. Recommended for large HPO runs. Defaults to True.	`True`
`distributed_trials`	`bool`	If False, trials run in the current process. If True and `trial_devices` resolves to more than one process, each trial is launched as a child `torch.distributed.run` job while the parent process remains the only Optuna controller.	`False`
`trial_devices`	`str \| int \| Sequence[int]`	Device/process selection for distributed trials. Accepted forms include `"auto"`, `"cpu"`, `"single"`, an int process count, a GPU-id sequence such as `[0, 1]`, or comma strings such as `"0,1"` / `"cuda:0,1"`.	`'auto'`
`parent_poll_interval_seconds`	`float`	Parent polling interval for distributed-trial IPC progress. Must be > 0.	`0.5`
`decision_wait_seconds`	`float`	Child-side wait time for a parent continue/prune decision after each reported epoch. Must be >= 0.	`20.0`
`child_exit_grace_seconds`	`float`	Grace period before the parent terminates a distributed child after a prune decision. Must be > 0.	`60.0`
`keep_trial_workdirs`	`bool`	Keep temporary distributed-trial workdirs for debugging. Defaults to False.	`False`
`sampler`	`OptunaSamplerConfig`	Sampler configuration.	`deepvisiontools.train.optuna_wrapper.OptunaSamplerConfig()`
`pruner`	`OptunaPrunerConfig`	Pruner configuration.	`deepvisiontools.train.optuna_wrapper.OptunaPrunerConfig()`
`score`	`OptunaScoreConfig`	Final objective reducer configuration.	`deepvisiontools.train.optuna_wrapper.OptunaScoreConfig()`
`plateau_stop`	`PlateauStopConfig`	Optional study-level plateau stopper.	`deepvisiontools.train.optuna_wrapper.PlateauStopConfig()`
`search_space`	`list[SearchSpaceSpec]`	Standard config-path search space.	`list()`
`augmentation_space`	`list[AugmentationPipelineSpec]`	Built-in augmentation optimization helper.	`list()`
`trial_mutator_target`	`Callable \| str \| None`	Optional function or import target for `fn(trial, cfg)`. Import strings are preferred for portable YAML configs.	`None`
`data_mutator_target`	`Callable \| str \| None`	Optional function or import target for `fn(trial, cfg)`. Use it for advanced dataset/loader rebuilding.	`None`

Minimal HPO

from deepvisiontools.train import TrainerConfig
from deepvisiontools.train.optuna_wrapper import (
    OptunaConfig,
    OptunaPrunerConfig,
    OptunaSamplerConfig,
    OptunaScoreConfig,
    SearchSpaceSpec,
    run_optuna,
)

cfg = TrainerConfig.build_registered_config(
    model="detector-seg",
    config="detector_convnextv2-base_cascade_maskrcnn",
    train_loader=trainloader,
    valid_loader=validloader,
)

cfg.epochs = 20
cfg.checkpoint.monitor = ("valid/loss", "loss")
cfg.checkpoint.mode = "min"

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="detector_hpo",
    storage="optuna_detector_hpo.db",
    load_if_exists=True,
    n_trials=30,
    sampler=OptunaSamplerConfig(
        kind="tpe",
        seed=42,
        n_startup_trials=8,
    ),
    pruner=OptunaPrunerConfig(
        kind="median",
        n_startup_trials=5,
        n_warmup_steps=3,
    ),
    score=OptunaScoreConfig(
        monitor=("valid/loss", "loss"),
        mode="min",
        reducer="best",
    ),
    search_space=[
        SearchSpaceSpec(
            name="lr",
            kind="float",
            path="optim.optim_args.lr",
            low=1e-5,
            high=1e-3,
            log=True,
        ),
        SearchSpaceSpec(
            name="weight_decay",
            kind="float",
            path="optim.weight_decay",
            low=1e-6,
            high=1e-2,
            log=True,
        ),
        SearchSpaceSpec(
            name="accumulate_steps",
            kind="categorical",
            path="precision.accumulate_steps",
            choices=[1, 2, 4],
        ),
    ],
)

study = run_optuna(cfg)
print(study.best_trial.number, study.best_value, study.best_params)

With augmentation search and plateau stopping

from deepvisiontools.train.optuna_wrapper import (
    AugmentationChoiceSpec,
    AugmentationPipelineSpec,
    AugmentationTransformSpec,
    OptunaConfig,
    PlateauStopConfig,
    SearchSpaceSpec,
    run_optuna,
)

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="augmentation_hpo",
    storage="augmentation_hpo.db",
    n_trials=100,
    plateau_stop=PlateauStopConfig(
        enabled=True,
        patience_trials=20,
        min_delta=1e-4,
    ),
    search_space=[
        SearchSpaceSpec(
            name="lr",
            kind="float",
            path="optim.optim_args.lr",
            low=1e-5,
            high=5e-4,
            log=True,
        ),
    ],
    augmentation_space=[
        AugmentationPipelineSpec(
            target_loaders=("train",),
            mode="replace",
            items=[
                AugmentationTransformSpec(
                    target="torchvision.transforms.v2:RandomHorizontalFlip",
                    kwargs={
                        "p": SearchSpaceSpec(
                            name="hflip_p",
                            kind="float",
                            low=0.0,
                            high=0.8,
                            step=0.1,
                        )
                    },
                    name="hflip",
                ),
                AugmentationChoiceSpec(
                    name="color_aug",
                    allow_none=True,
                    choices=[
                        AugmentationTransformSpec(
                            target="torchvision.transforms.v2:ColorJitter",
                            kwargs={
                                "brightness": 0.15,
                                "contrast": 0.15,
                                "saturation": 0.15,
                                "hue": 0.05,
                            },
                            name="color_jitter",
                        ),
                        AugmentationTransformSpec(
                            target="torchvision.transforms.v2:RandomGrayscale",
                            kwargs={"p": 0.05},
                            name="grayscale",
                        ),
                    ],
                ),
            ],
        )
    ],
)

study = run_optuna(cfg)

Distributed trial execution

from deepvisiontools.train.optuna_wrapper import OptunaConfig, run_optuna

cfg.optuna = OptunaConfig(
    enabled=True,
    study_name="ddp_trial_hpo",
    storage="ddp_trial_hpo.db",
    n_trials=20,

    # One Optuna trial at a time, but each trial is a DDP subprocess job.
    distributed_trials=True,
    trial_devices=[0, 1],
    parent_poll_interval_seconds=0.5,
    decision_wait_seconds=20.0,
    child_exit_grace_seconds=60.0,

    # Useful while debugging subprocess trial payloads / IPC files.
    keep_trial_workdirs=False,
)

study = run_optuna(cfg)

`OptunaSamplerConfig(kind='tpe', seed=None, n_startup_trials=10, multivariate=True, group=True)` `dataclass` ¶

Optuna sampler configuration.

Parameters:

Name	Type	Description	Default
`kind`	`Literal["tpe", "random", "cmaes"]`	Sampler kind. `"tpe"` builds `optuna.samplers.TPESampler`, `"random"` builds `optuna.samplers.RandomSampler`, and `"cmaes"` builds `optuna.samplers.CmaEsSampler`. Defaults to `"tpe"`.	`'tpe'`
`seed`	`int \| None`	Optional random seed forwarded to the sampler.	`None`
`n_startup_trials`	`int`	Number of initial random/exploration trials for TPE and CMA-ES samplers. Defaults to 10.	`10`
`multivariate`	`bool`	Forwarded to `TPESampler` only. Enables multivariate TPE. Defaults to True.	`True`
`group`	`bool`	Forwarded to `TPESampler` only. Enables grouped multivariate search. Defaults to True.	`True`

Usage

from deepvisiontools.train.optuna_wrapper import OptunaSamplerConfig

# Default/recommended sampler.
sampler = OptunaSamplerConfig(
    kind="tpe",
    seed=42,
    n_startup_trials=10,
    multivariate=True,
    group=True,
)

# Simpler reproducible baseline.
random_sampler = OptunaSamplerConfig(
    kind="random",
    seed=42,
)

`OptunaPrunerConfig(kind='median', n_startup_trials=5, n_warmup_steps=0, interval_steps=1, min_resource=1, reduction_factor=3, max_resource='auto')` `dataclass` ¶

Optuna pruner configuration.

Pruning is driven by the monitored value resolved after each epoch. In local trials this is handled by OptunaPruningHook. In distributed trials, the child process writes progress to IPC files and the parent Optuna controller makes pruning decisions.

Parameters:

Name	Type	Description	Default
`kind`	`Literal["median", "successive_halving", "hyperband", "none"]`	Pruner kind. Use `"none"` to disable pruning. Defaults to `"median"`.	`'median'`
`n_startup_trials`	`int`	Number of completed trials before the median pruner can prune. Used by `"median"`.	`5`
`n_warmup_steps`	`int`	Number of initial reported steps ignored by the median pruner. Since reports are epoch-based here, this is usually interpreted as warmup epochs. Used by `"median"`.	`0`
`interval_steps`	`int`	Pruning check interval. Used by `"median"`.	`1`
`min_resource`	`int`	Minimum resource for successive halving or hyperband. In this training integration, resource corresponds to reported epoch steps.	`1`
`reduction_factor`	`int`	Reduction factor for successive halving or hyperband.	`3`
`max_resource`	`str \| int`	Maximum resource for hyperband. `"auto"` lets Optuna infer it. Otherwise provide an int.	`'auto'`

Usage

from deepvisiontools.train.optuna_wrapper import OptunaPrunerConfig

# Median pruning, ignoring the first 3 reported epochs.
pruner = OptunaPrunerConfig(
    kind="median",
    n_startup_trials=5,
    n_warmup_steps=3,
    interval_steps=1,
)

# Disable pruning completely.
no_pruning = OptunaPrunerConfig(kind="none")

# More aggressive resource-based pruning.
sha = OptunaPrunerConfig(
    kind="successive_halving",
    min_resource=1,
    reduction_factor=3,
)

`OptunaScoreConfig(monitor=None, mode=None, reducer='best', top_k=3, callback_target=None)` `dataclass` ¶

Final per-trial objective score configuration.

During training, the monitored value is recorded after each epoch. At the end of the trial, this history is reduced into one scalar objective returned to Optuna.

Parameters:

Name	Type	Description	Default
`monitor`	`str \| tuple[str, ...] \| list[str] \| None`	Metric/loss path to monitor. If None, falls back to `cfg.checkpoint.monitor`. Examples: `("valid/loss", "loss")` for validation loss or `"valid/map"` for a flat metric key.	`None`
`mode`	`Literal["min", "max"] \| None`	Direction of the objective. If None, falls back to `cfg.checkpoint.mode`. Use `"min"` for losses and `"max"` for metrics.	`None`
`reducer`	`Literal["best", "last", "top_k_mean"]`	How to convert the per-epoch history into the final objective. `"best"` selects the best value according to `mode`; `"last"` uses the final recorded value; `"top_k_mean"` averages the best `top_k` values.	`'best'`
`top_k`	`int`	Number of best values used by `"top_k_mean"`. Defaults to 3.	`3`
`callback_target`	`str \| None`	Optional import target for a custom score reducer. The callable receives `fn(history, *, mode, cfg, trial, trainer)` and must return either a float objective or a dict containing `"objective"` and optional `"stats"`.	`None`

Usage

from deepvisiontools.train.optuna_wrapper import OptunaScoreConfig

# Optimize the best validation loss observed during the trial.
score = OptunaScoreConfig(
    monitor=("valid/loss", "loss"),
    mode="min",
    reducer="best",
)

# More stable score: average the 3 best validation mAP values.
map_score = OptunaScoreConfig(
    monitor="valid/map",
    mode="max",
    reducer="top_k_mean",
    top_k=3,
)

# Custom reducer.
custom_score = OptunaScoreConfig(
    monitor=("valid/loss", "loss"),
    mode="min",
    reducer="best",
    callback_target="my_project.hpo_callbacks:score_trial",
)

Custom callback

# my_project/hpo_callbacks.py

def score_trial(history, *, mode, cfg, trial, trainer):
    values = [float(row["value"]) for row in history]
    objective = min(values) if mode == "min" else max(values)

    return {
        "objective": objective,
        "stats": {
            "num_epochs": len(values),
            "last_value": values[-1],
            "best_value": objective,
        },
    }

`PlateauStopConfig(enabled=False, patience_trials=20, min_delta=0.0)` `dataclass` ¶

Study-level plateau stopping configuration.

This stops the whole Optuna study when the best completed trial has not improved enough over a recent window of completed trials.

Parameters:

Name	Type	Description	Default
`enabled`	`bool`	Enable study-level plateau stopping. Defaults to False.	`False`
`patience_trials`	`int`	Number of completed trials to wait without significant improvement before stopping. Defaults to 20.	`20`
`min_delta`	`float`	Minimum improvement required to reset patience. For `"min"`, improvement means the best value decreased by more than `min_delta`; for `"max"`, it increased by more than `min_delta`.	`0.0`

Usage

from deepvisiontools.train.optuna_wrapper import PlateauStopConfig

plateau_stop = PlateauStopConfig(
    enabled=True,
    patience_trials=15,
    min_delta=1e-4,
)

`SearchSpaceSpec(name, kind, path=None, low=None, high=None, step=None, log=False, choices=None)` `dataclass` ¶

One Optuna suggestion, optionally written into a TrainerConfig path.

SearchSpaceSpec is the basic building block for standard hyperparameter search. During each trial, suggest(trial) calls the corresponding Optuna trial.suggest_* method. If path is provided, OptunaRunner writes the suggested value into the cloned trial config before creating the Trainer.

Parameters:

Name	Type	Description	Default
`name`	`str`	Optuna parameter name shown in the study, dashboard, and trial params.	required
`kind`	`Literal["float", "int", "categorical", "bool"]`	Suggestion type. `"float"` uses `trial.suggest_float`, `"int"` uses `trial.suggest_int`, `"categorical"` uses `trial.suggest_categorical`, and `"bool"` is implemented as a categorical choice between `False` and `True`.	required
`path`	`str \| None`	Optional dotted path in the trial `TrainerConfig` where the suggested value is written. Supports dataclass/object attributes, dictionaries, and list indices. Examples: `"optim.optim_args.lr"`, `"precision.accumulate_steps"`, `"model_spec.kwargs.arch"`. If None, the value is suggested but not automatically assigned, which is useful when using custom mutators.	`None`
`low`	`float \| None`	Lower bound for `"float"` or `"int"` suggestions.	`None`
`high`	`float \| None`	Upper bound for `"float"` or `"int"` suggestions.	`None`
`step`	`float \| None`	Optional step size. For `"float"`, setting `step` disables logarithmic sampling. For `"int"`, it is converted to int.	`None`
`log`	`bool`	Whether to use log-scale sampling for `"float"` or `"int"`. For `"int"`, log sampling is only used when `step is None`.	`False`
`choices`	`list[Any] \| None`	Candidate values for `"categorical"`.	`None`

Usage

from deepvisiontools.train.optuna_wrapper import SearchSpaceSpec

search_space = [
    # Optimizer learning rate.
    SearchSpaceSpec(
        name="lr",
        kind="float",
        path="optim.optim_args.lr",
        low=1e-5,
        high=1e-3,
        log=True,
    ),

    # Weight decay.
    SearchSpaceSpec(
        name="weight_decay",
        kind="float",
        path="optim.weight_decay",
        low=1e-6,
        high=1e-2,
        log=True,
    ),

    # Gradient accumulation.
    SearchSpaceSpec(
        name="accumulate_steps",
        kind="categorical",
        path="precision.accumulate_steps",
        choices=[1, 2, 4],
    ),

    # Enable / disable EMA.
    SearchSpaceSpec(
        name="use_ema",
        kind="bool",
        path="ema.enabled",
    ),

    # Tune a model builder argument stored in ModelSpec.
    SearchSpaceSpec(
        name="backbone",
        kind="categorical",
        path="model_spec.kwargs.backbone_name",
        choices=["convnextv2_base", "convnextv2_large"],
    ),
]

`AugmentationTransformSpec(target, kwargs=dict(), enabled=True, name=None)` `dataclass` ¶

One importable augmentation transform candidate for Optuna trials.

The transform is instantiated during each trial. Constant keyword arguments are copied as-is. Keyword arguments that are SearchSpaceSpec objects are suggested per trial before the transform is built.

Parameters:

Name	Type	Description	Default
`target`	`str`	Import target of the transform class or callable, using `"module:qualname"` or `"module.qualname"` syntax. Typical values are torchvision v2 transforms such as `"torchvision.transforms.v2:RandomHorizontalFlip"`.	required
`kwargs`	`dict[str, Any]`	Keyword arguments passed to the transform constructor. Values can be constants or `SearchSpaceSpec` instances.	`dict()`
`enabled`	`bool \| SearchSpaceSpec`	If bool, always enable or disable the transform. If a `SearchSpaceSpec` is provided, it is suggested per trial and converted to bool.	`True`
`name`	`str \| None`	Optional human-readable name used by `AugmentationChoiceSpec`. If None, `target` is used.	`None`

Usage

from deepvisiontools.train.optuna_wrapper import (
    AugmentationTransformSpec,
    SearchSpaceSpec,
)

horizontal_flip = AugmentationTransformSpec(
    target="torchvision.transforms.v2:RandomHorizontalFlip",
    kwargs={
        "p": SearchSpaceSpec(
            name="hflip_p",
            kind="float",
            low=0.0,
            high=0.8,
            step=0.1,
        ),
    },
    enabled=True,
    name="horizontal_flip",
)

optional_grayscale = AugmentationTransformSpec(
    target="torchvision.transforms.v2:RandomGrayscale",
    kwargs={"p": 0.05},
    enabled=SearchSpaceSpec(
        name="enable_grayscale",
        kind="bool",
    ),
    name="grayscale",
)

`AugmentationChoiceSpec(name, choices=list(), allow_none=False)` `dataclass` ¶

Categorical choice among several augmentation transform candidates.

Use this when exactly one augmentation family should be selected during a trial, for example choosing one color augmentation strategy. If allow_none is True, Optuna can also choose to apply no transform for this choice.

Parameters:

Name	Type	Description	Default
`name`	`str`	Optuna categorical parameter name used for the choice.	required
`choices`	`list[AugmentationTransformSpec]`	Candidate transforms.	`list()`
`allow_none`	`bool`	If True, adds a `"__none__"` choice meaning no transform is materialized for this item.	`False`

Usage

from deepvisiontools.train.optuna_wrapper import (
    AugmentationChoiceSpec,
    AugmentationTransformSpec,
    SearchSpaceSpec,
)

color_choice = AugmentationChoiceSpec(
    name="color_aug",
    allow_none=True,
    choices=[
        AugmentationTransformSpec(
            target="torchvision.transforms.v2:ColorJitter",
            kwargs={
                "brightness": SearchSpaceSpec(
                    name="jitter_brightness",
                    kind="float",
                    low=0.0,
                    high=0.3,
                    step=0.05,
                ),
                "contrast": 0.15,
                "saturation": 0.15,
                "hue": 0.05,
            },
            name="color_jitter",
        ),
        AugmentationTransformSpec(
            target="torchvision.transforms.v2:RandomAutocontrast",
            kwargs={"p": 0.5},
            name="autocontrast",
        ),
    ],
)

`AugmentationPipelineSpec(target_loaders=('train',), mode='replace', items=list())` `dataclass` ¶

Trial-time augmentation pipeline mutation helper.

For each trial, OptunaRunner materializes the configured transform items, deep-copies the target loader dataset, replaces or appends to its dataset.augmentation list, then rebuilds a loader of the same class when possible.

Parameters:

Name	Type	Description	Default
`target_loaders`	`tuple[str, ...]`	Loader names to mutate. Supported names follow `TrainerConfig` attributes: `"train"` targets `cfg.train_loader`, and `"valid"` targets `cfg.valid_loader`. Defaults to `("train",)`.	`('train',)`
`mode`	`Literal["replace", "append"]`	If `"replace"`, the dataset augmentation list is replaced by the trial transforms. If `"append"`, trial transforms are appended to the existing augmentation list.	`'replace'`
`items`	`list[AugmentationTransformSpec \| AugmentationChoiceSpec]`	Ordered transform specs or categorical transform choices.	`list()`

Usage

from deepvisiontools.train.optuna_wrapper import (
    AugmentationChoiceSpec,
    AugmentationPipelineSpec,
    AugmentationTransformSpec,
    SearchSpaceSpec,
)

augmentation_space = [
    AugmentationPipelineSpec(
        target_loaders=("train",),
        mode="replace",
        items=[
            AugmentationTransformSpec(
                target="torchvision.transforms.v2:RandomHorizontalFlip",
                kwargs={
                    "p": SearchSpaceSpec(
                        name="hflip_p",
                        kind="float",
                        low=0.0,
                        high=0.8,
                        step=0.1,
                    )
                },
                name="hflip",
            ),
            AugmentationChoiceSpec(
                name="geometric_aug",
                allow_none=True,
                choices=[
                    AugmentationTransformSpec(
                        target="torchvision.transforms.v2:RandomVerticalFlip",
                        kwargs={"p": 0.5},
                        name="vflip",
                    ),
                    AugmentationTransformSpec(
                        target="torchvision.transforms.v2:RandomRotation",
                        kwargs={"degrees": 15},
                        name="rotation15",
                    ),
                ],
            ),
        ],
    )
]

`OptunaPruningHook(monitor, *, trial_attr_name='optuna_trial', min_step=0)` ¶

Bases: deepvisiontools.train.hooks.Hook

In-process pruning hook used for the single-process path.

This hook is not used for distributed subprocess trials. In the DDP trial path, the parent Optuna controller owns the Trial and the child communicates through IPC instead.

Source code in src/deepvisiontools/train/optuna_wrapper.py

def __init__(
    self,
    monitor: str | tuple[str, ...] | list[str],
    *,
    trial_attr_name: str = "optuna_trial",
    min_step: int = 0,
):
    self.monitor = monitor
    self.trial_attr_name = str(trial_attr_name)
    self.min_step = max(int(min_step), 0)

`OptunaIPCProgressHook(monitor, *, ipc_dir, decision_wait_seconds=20.0, decision_poll_interval_seconds=0.25)` ¶

Bases: deepvisiontools.train.hooks.Hook

Child-process hook used by distributed subprocess trials.

Behavior¶

rank 0 resolves the monitored value after each epoch
rank 0 appends one JSONL progress record
rank 0 waits briefly for a parent decision for that step:
- continue
- prune
pruning intent is synchronized across ranks with _sync_any_true(...)

Why this design¶

The parent process owns the real Optuna Trial object. The child never talks to Optuna directly; it only emits progress and obeys decisions.

Source code in src/deepvisiontools/train/optuna_wrapper.py

def __init__(
    self,
    monitor: str | tuple[str, ...] | list[str],
    *,
    ipc_dir: str | Path,
    decision_wait_seconds: float = 20.0,
    decision_poll_interval_seconds: float = 0.25,
):
    self.monitor = monitor
    self.ipc_dir = Path(ipc_dir)
    self.decision_wait_seconds = max(float(decision_wait_seconds), 0.0)
    self.decision_poll_interval_seconds = max(
        float(decision_poll_interval_seconds), 0.05
    )

    self.progress_path = self.ipc_dir / "progress.jsonl"
    self.decision_path = self.ipc_dir / "decision.json"

Optuna HPO

run_optuna(cfg) ¶

OptunaRunner(base_cfg) ¶

OptunaSamplerConfig(kind='tpe', seed=None, n_startup_trials=10, multivariate=True, group=True) dataclass ¶

OptunaPrunerConfig(kind='median', n_startup_trials=5, n_warmup_steps=0, interval_steps=1, min_resource=1, reduction_factor=3, max_resource='auto') dataclass ¶

OptunaScoreConfig(monitor=None, mode=None, reducer='best', top_k=3, callback_target=None) dataclass ¶

PlateauStopConfig(enabled=False, patience_trials=20, min_delta=0.0) dataclass ¶

SearchSpaceSpec(name, kind, path=None, low=None, high=None, step=None, log=False, choices=None) dataclass ¶

AugmentationTransformSpec(target, kwargs=dict(), enabled=True, name=None) dataclass ¶

AugmentationChoiceSpec(name, choices=list(), allow_none=False) dataclass ¶

AugmentationPipelineSpec(target_loaders=('train',), mode='replace', items=list()) dataclass ¶

OptunaPruningHook(monitor, *, trial_attr_name='optuna_trial', min_step=0) ¶

OptunaIPCProgressHook(monitor, *, ipc_dir, decision_wait_seconds=20.0, decision_poll_interval_seconds=0.25) ¶

Behavior¶

Why this design¶

`run_optuna(cfg)` ¶

`OptunaRunner(base_cfg)` ¶

`OptunaSamplerConfig(kind='tpe', seed=None, n_startup_trials=10, multivariate=True, group=True)` `dataclass` ¶

`OptunaPrunerConfig(kind='median', n_startup_trials=5, n_warmup_steps=0, interval_steps=1, min_resource=1, reduction_factor=3, max_resource='auto')` `dataclass` ¶

`OptunaScoreConfig(monitor=None, mode=None, reducer='best', top_k=3, callback_target=None)` `dataclass` ¶

`PlateauStopConfig(enabled=False, patience_trials=20, min_delta=0.0)` `dataclass` ¶

`SearchSpaceSpec(name, kind, path=None, low=None, high=None, step=None, log=False, choices=None)` `dataclass` ¶

`AugmentationTransformSpec(target, kwargs=dict(), enabled=True, name=None)` `dataclass` ¶

`AugmentationChoiceSpec(name, choices=list(), allow_none=False)` `dataclass` ¶

`AugmentationPipelineSpec(target_loaders=('train',), mode='replace', items=list())` `dataclass` ¶

`OptunaPruningHook(monitor, *, trial_attr_name='optuna_trial', min_step=0)` ¶

`OptunaIPCProgressHook(monitor, *, ipc_dir, decision_wait_seconds=20.0, decision_poll_interval_seconds=0.25)` ¶