objective

orchard.optimization.objective.objective

Optuna Objective Function for Vision Pipeline.

Provides OptunaObjective, the callable that Optuna invokes for each trial. It orchestrates the full trial lifecycle — config sampling, data loading, model creation, training, and metric extraction — using dependency injection for maximum testability.

Key Components:

• OptunaObjective: High-level orchestration callable. All external dependencies (dataset loader, dataloader factory, model factory) are injectable via Protocol-based abstractions. Settings are read from cfg.optuna.* as single source of truth.
• TrialConfigBuilder: Builds trial-specific Config instances from sampled hyperparameters.
• MetricExtractor: Handles metric extraction and best-value tracking across epochs.
• TrialTrainingExecutor: Executes training loops with Optuna pruning integration.

Example

objective = OptunaObjective(cfg, search_space, device) study = optuna.create_study(direction="maximize") study.optimize(objective, n_trials=50)

DatasetLoaderProtocol

Bases: Protocol

Protocol for dataset loading (enables dependency injection).

__call__(metadata)

Load dataset from metadata.

Source code in orchard/optimization/objective/objective.py

def __call__(self, metadata: DatasetMetadata) -> DatasetData:
    """Load dataset from metadata."""
    ...  # pragma: no cover

DataloaderFactoryProtocol

Bases: Protocol

Protocol for dataloader creation (enables dependency injection).

__call__(metadata, dataset_cfg, training_cfg, aug_cfg, num_workers, is_optuna=False)

Create train/val/test dataloaders.

Source code in orchard/optimization/objective/objective.py

def __call__(
    self,
    metadata: DatasetData,
    dataset_cfg: DatasetConfig,
    training_cfg: TrainingConfig,
    aug_cfg: AugmentationConfig,
    num_workers: int,
    is_optuna: bool = False,
) -> tuple[DataLoader[Any], DataLoader[Any], DataLoader[Any]]:
    """Create train/val/test dataloaders."""
    ...  # pragma: no cover

ModelFactoryProtocol

Bases: Protocol

Protocol for model creation (enables dependency injection).

__call__(device, dataset_cfg, arch_cfg)

Create and initialize model.

Source code in orchard/optimization/objective/objective.py

def __call__(
    self,
    device: torch.device,
    dataset_cfg: DatasetConfig,
    arch_cfg: ArchitectureConfig,
) -> torch.nn.Module:
    """Create and initialize model."""
    ...  # pragma: no cover

OptunaObjective(cfg, search_space, device, dataset_loader=None, dataloader_factory=None, model_factory=None, tracker=None)

Optuna objective function with dependency injection.

Orchestrates hyperparameter optimization trials by:

• Building trial-specific configurations
• Creating data loaders, models, and optimizers
• Executing training with pruning
• Tracking and returning best metrics

All external dependencies are injectable for testability:

• dataset_loader: Dataset loading function
• dataloader_factory: DataLoader creation function
• model_factory: Model instantiation function

Attributes:

Name	Type	Description
cfg		Base configuration (single source of truth)
search_space		Hyperparameter search space
device		Training device (CPU/CUDA/MPS)
config_builder		Builds trial-specific configs
metric_extractor		Handles metric extraction
dataset_data		Cached dataset (loaded once, reused across trials)

Example

objective = OptunaObjective( ... cfg=config, ... search_space=search_space, ... device=torch.device("cuda"), ... ) study = optuna.create_study(direction="maximize") study.optimize(objective, n_trials=50)

Initialize Optuna objective.

Parameters:

Name	Type	Description	Default
cfg	Config	Base configuration (reads optuna.* settings)	required
search_space	Mapping[str, Any]	Hyperparameter search space	required
device	device	Training device	required
dataset_loader	DatasetLoaderProtocol | None	Dataset loading function (default: load_dataset)	None
dataloader_factory	DataloaderFactoryProtocol | None	DataLoader factory (default: get_dataloaders)	None
model_factory	ModelFactoryProtocol | None	Model factory (default: get_model)	None
tracker	TrackerProtocol | None	Optional experiment tracker for nested trial logging	None

Source code in orchard/optimization/objective/objective.py

def __init__(
    self,
    cfg: Config,
    search_space: Mapping[str, Any],
    device: torch.device,
    dataset_loader: DatasetLoaderProtocol | None = None,
    dataloader_factory: DataloaderFactoryProtocol | None = None,
    model_factory: ModelFactoryProtocol | None = None,
    tracker: TrackerProtocol | None = None,
) -> None:
    """
    Initialize Optuna objective.

    Args:
        cfg: Base configuration (reads optuna.* settings)
        search_space: Hyperparameter search space
        device: Training device
        dataset_loader: Dataset loading function (default: load_dataset)
        dataloader_factory: DataLoader factory (default: get_dataloaders)
        model_factory: Model factory (default: get_model)
        tracker: Optional experiment tracker for nested trial logging
    """
    self.cfg = cfg
    self.search_space = search_space
    self.device = device
    self.tracker = tracker

    # Dependency injection with defaults
    self._dataset_loader = dataset_loader or load_dataset
    self._dataloader_factory = dataloader_factory or get_dataloaders
    self._model_factory = model_factory or get_model

    # Components (monitor_metric is the single source of truth for the
    # optimisation target — shared by trainer checkpointing and Optuna ranking)
    self.config_builder = TrialConfigBuilder(cfg)
    self.metric_extractor = MetricExtractor(
        cfg.training.monitor_metric, direction=cfg.optuna.direction
    )

    # Load dataset once (reused across all trials)
    self.dataset_data = self._dataset_loader(self.config_builder.base_metadata)

__call__(trial)

Execute single Optuna trial.

Samples hyperparameters, builds trial configuration, trains model, and returns best validation metric. Failed trials return the worst possible metric instead of crashing the study.

Parameters:

Name	Type	Description	Default
trial	Trial	Optuna trial object	required

Returns:

Type	Description
float	Best validation metric achieved during training,
float	or worst-case metric if the trial fails.

Raises:

Type	Description
TrialPruned	If trial is pruned during training

Source code in orchard/optimization/objective/objective.py

def __call__(self, trial: optuna.Trial) -> float:
    """
    Execute single Optuna trial.

    Samples hyperparameters, builds trial configuration, trains model,
    and returns best validation metric. Failed trials return the worst
    possible metric instead of crashing the study.

    Args:
        trial: Optuna trial object

    Returns:
        Best validation metric achieved during training,
        or worst-case metric if the trial fails.

    Raises:
        optuna.TrialPruned: If trial is pruned during training
    """
    # Reset per-trial metric tracking
    self.metric_extractor.reset()

    # Sample parameters
    params = self._sample_params(trial)

    # Build trial config
    trial_cfg = self.config_builder.build(params)

    # Inject recipe-level flags for logging (not Optuna params)
    log_params = {**params, "pretrained": self.cfg.architecture.pretrained}

    # Log trial start
    log_trial_start(trial.number, log_params)

    # Start nested MLflow run for this trial
    if self.tracker is not None:
        self.tracker.start_optuna_trial(trial.number, log_params)

    trial_succeeded = False
    try:
        # Setup training components
        train_loader, val_loader, _ = self._dataloader_factory(
            self.dataset_data,
            trial_cfg.dataset,
            trial_cfg.training,
            trial_cfg.augmentation,
            trial_cfg.num_workers,
            is_optuna=True,
        )
        model = self._model_factory(self.device, trial_cfg.dataset, trial_cfg.architecture)
        optimizer = get_optimizer(model, trial_cfg.training)
        scheduler = get_scheduler(optimizer, trial_cfg.training)

        class_weights = None
        if trial_cfg.training.weighted_loss:
            train_labels = train_loader.dataset.labels.flatten()  # type: ignore[attr-defined]
            num_classes = self.config_builder.base_metadata.num_classes
            class_weights = compute_class_weights(train_labels, num_classes, self.device)

        criterion = get_criterion(trial_cfg.training, class_weights=class_weights)

        # Execute training
        executor = TrialTrainingExecutor(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            optimizer=optimizer,
            scheduler=scheduler,
            criterion=criterion,
            training=trial_cfg.training,
            optuna=trial_cfg.optuna,
            log_interval=trial_cfg.telemetry.log_interval,
            device=self.device,
            metric_extractor=self.metric_extractor,
        )

        best_metric = executor.execute(trial)
        trial_succeeded = True

        return best_metric

    except optuna.TrialPruned:
        trial_succeeded = True  # pruned trials have valid metrics
        raise

    except Exception as e:  # must not crash study
        logger.error(  # pragma: no mutate
            "%s%s Trial %d failed: %s: %s",
            LogStyle.INDENT,  # pragma: no mutate
            LogStyle.FAILURE,  # pragma: no mutate
            trial.number,  # pragma: no mutate
            type(e).__name__,  # pragma: no mutate
            e,  # pragma: no mutate
        )
        return self._worst_metric()

    finally:
        # End nested MLflow run for this trial
        if self.tracker is not None:
            if trial_succeeded:
                self.tracker.end_optuna_trial(self.metric_extractor.best_metric)
            else:
                # Trial failed before any validation — close run without metric
                self.tracker.end_optuna_trial(self._worst_metric())

        # Cleanup GPU memory between trials
        self._cleanup()