|
| 1 | +"""Unit that measures a fold run with a search of its own inside each fold.""" |
| 2 | + |
| 3 | +import logging |
| 4 | + |
| 5 | +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum |
| 6 | +from DashAI.back.core.schema_fields import BaseSchema |
| 7 | +from DashAI.back.job.base_job import JobError |
| 8 | +from DashAI.back.units.context import ExecutionContext |
| 9 | +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit |
| 10 | +from DashAI.back.units.fit_scope import ( |
| 11 | + TRIAL_SPLITS, |
| 12 | + goal_metric_field, |
| 13 | + optimizer_field, |
| 14 | + trial_splits_field, |
| 15 | +) |
| 16 | +from DashAI.back.units.splitter_scope import _splitter_field |
| 17 | + |
| 18 | +log = logging.getLogger(__name__) |
| 19 | + |
| 20 | + |
| 21 | +def inner_splitter_field(): |
| 22 | + """The splitter that carves an outer fold into folds of its own. |
| 23 | +
|
| 24 | + Required rather than optional, which is what makes this a separate unit. |
| 25 | + A component field wrapped in ``none_type`` is emitted as ``anyOf`` and the |
| 26 | + front reads ``parent`` directly off the property, so an optional one leaves |
| 27 | + the user with no selector at all -- the same wall that made the two |
| 28 | + explainer units siblings. |
| 29 | + """ |
| 30 | + return _splitter_field( |
| 31 | + parent="FoldSplitter", |
| 32 | + placeholder={ |
| 33 | + "component": "KFoldSplitter", |
| 34 | + "params": {"n_splits": 3, "shuffle": True, "random_state": 42}, |
| 35 | + }, |
| 36 | + ) |
| 37 | + |
| 38 | + |
| 39 | +class FitModelOverNestedFoldsSchema(BaseSchema): |
| 40 | + optimizer: optimizer_field() # type: ignore |
| 41 | + goal_metric: goal_metric_field() # type: ignore |
| 42 | + scored_splits: trial_splits_field() # type: ignore |
| 43 | + inner_splitter: inner_splitter_field() # type: ignore |
| 44 | + |
| 45 | + |
| 46 | +class FitModelOverNestedFoldsUnit(FitModelOverFoldsUnit): |
| 47 | + """Score every outer fold with a search that never saw it. |
| 48 | +
|
| 49 | + ``FitModelOverFoldsUnit`` plus one measurement taken before it. The |
| 50 | + relationship is inheritance rather than a shared mixin because that is what |
| 51 | + it is: everything the sibling does still happens, and this adds a step in |
| 52 | + front of it. |
| 53 | +
|
| 54 | + **What the extra step is for.** In an ordinary cross-validated search, the |
| 55 | + same folds choose the hyperparameters and report the score, so the reported |
| 56 | + score is optimistic by however much the search managed to fit them. The |
| 57 | + nested loop measures that honestly: for each outer fold, a search is run |
| 58 | + from scratch on folds carved out of *that fold's training rows only*, and |
| 59 | + the chosen model is then scored on the outer fold's validation rows, which |
| 60 | + that search never saw. Those are the ``OUTER_FOLD`` rows. |
| 61 | +
|
| 62 | + What it does **not** do is choose the hyperparameters. Each outer fold |
| 63 | + picks its own, and they generally differ; there is no single model to keep |
| 64 | + out of that loop. So the ordinary search still runs afterwards, over all |
| 65 | + the folds, and it is what produces the model that gets saved. The nested |
| 66 | + numbers are a statement about the procedure, not about the artifact. |
| 67 | +
|
| 68 | + That also means the inner trials record nothing: their partitions belong to |
| 69 | + one outer fold, and rows from them would sit alongside rows describing the |
| 70 | + kept model as if they were comparable. |
| 71 | + """ |
| 72 | + |
| 73 | + SCHEMA = FitModelOverNestedFoldsSchema |
| 74 | + |
| 75 | + PROVIDES = ("model", "plot_paths", "fold_metrics", "outer_fold_metrics") |
| 76 | + |
| 77 | + def __init__(self, **config) -> None: |
| 78 | + super().__init__(**config) |
| 79 | + self._inner_splitter = None |
| 80 | + |
| 81 | + def _resolve_inner_splitter(self): |
| 82 | + """Build the splitter that carves an outer fold, memoized on this unit.""" |
| 83 | + if self._inner_splitter is not None: |
| 84 | + return self._inner_splitter |
| 85 | + |
| 86 | + from kink import di |
| 87 | + |
| 88 | + chosen = self.config["inner_splitter"] |
| 89 | + try: |
| 90 | + splitter_class = di["component_registry"][chosen["component"]]["class"] |
| 91 | + self._inner_splitter = splitter_class(splits_data=dict(chosen["params"])) |
| 92 | + except Exception as e: |
| 93 | + log.exception(e) |
| 94 | + raise JobError( |
| 95 | + f"Error configuring inner splitter for nested CV: {e}", |
| 96 | + ) from e |
| 97 | + return self._inner_splitter |
| 98 | + |
| 99 | + def execute(self, ctx: ExecutionContext) -> None: |
| 100 | + model = ctx.require("model") |
| 101 | + x_folds = ctx.require("x_folds") |
| 102 | + y_folds = ctx.require("y_folds") |
| 103 | + optimizable_parameters = ctx.require("optimizable_parameters") |
| 104 | + |
| 105 | + # Resolved outside the wrapper below so a splitter that cannot be built |
| 106 | + # is reported as that, rather than as a training failure. |
| 107 | + inner_splitter = self._resolve_inner_splitter() |
| 108 | + |
| 109 | + if optimizable_parameters: |
| 110 | + try: |
| 111 | + outer_fold_metrics = self._measure_every_outer_fold( |
| 112 | + model, |
| 113 | + x_folds, |
| 114 | + y_folds, |
| 115 | + inner_splitter, |
| 116 | + optimizable_parameters, |
| 117 | + ) |
| 118 | + except Exception as e: |
| 119 | + log.exception(e) |
| 120 | + raise JobError( |
| 121 | + f"Model training failed {e}", |
| 122 | + ) from e |
| 123 | + else: |
| 124 | + # Nothing to search means nothing for the nested loop to measure: |
| 125 | + # every outer fold would choose the same parameters, which is what |
| 126 | + # the ordinary loop already reports. |
| 127 | + outer_fold_metrics = {} |
| 128 | + |
| 129 | + ctx.put_ref("outer_fold_metrics", outer_fold_metrics) |
| 130 | + |
| 131 | + # And then the ordinary run, which is what produces the kept model. |
| 132 | + super().execute(ctx) |
| 133 | + |
| 134 | + def _measure_every_outer_fold( |
| 135 | + self, model, x_folds, y_folds, inner_splitter, optimizable_parameters |
| 136 | + ) -> dict: |
| 137 | + """Search inside each outer fold, then score that fold with what it chose. |
| 138 | +
|
| 139 | + Returns |
| 140 | + ------- |
| 141 | + dict |
| 142 | + ``{split name: {metric name: [one score per outer fold]}}``, for |
| 143 | + whoever aggregates them -- the same shape and the same reason as |
| 144 | + the ordinary fold scores. |
| 145 | + """ |
| 146 | + optimizer, goal_metric = self._resolve_search() |
| 147 | + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) |
| 148 | + outer_fold_metrics: dict = {} |
| 149 | + |
| 150 | + for index, (x_outer, y_outer) in enumerate(self._folds(x_folds, y_folds)): |
| 151 | + # Carved out of this fold's training rows alone. Its validation |
| 152 | + # rows are what the search will be judged on, so nothing drawn from |
| 153 | + # them may reach it. |
| 154 | + inner_x, inner_y, _ = inner_splitter.split( |
| 155 | + x_outer["train"], y_outer["train"] |
| 156 | + ) |
| 157 | + |
| 158 | + optimizer.optimize( |
| 159 | + model, |
| 160 | + inner_x, |
| 161 | + inner_y, |
| 162 | + optimizable_parameters, |
| 163 | + goal_metric, |
| 164 | + self._score_one_inner_trial, |
| 165 | + ) |
| 166 | + outer_model = optimizer.get_model() |
| 167 | + |
| 168 | + outer_model.x_data = x_outer |
| 169 | + outer_model.y_data = y_outer |
| 170 | + self._fit_kept_model(outer_model, x_outer, y_outer) |
| 171 | + |
| 172 | + for name in scored_splits: |
| 173 | + split = SplitEnum[name] |
| 174 | + results = outer_model.calculate_metrics( |
| 175 | + split=split, level=LevelEnum.OUTER_FOLD, fold_index=index |
| 176 | + ) |
| 177 | + if results is None: |
| 178 | + results = outer_model.compute_metrics(split=split) |
| 179 | + if not results: |
| 180 | + continue |
| 181 | + for metric_name, value in results.items(): |
| 182 | + outer_fold_metrics.setdefault(name, {}).setdefault( |
| 183 | + metric_name, [] |
| 184 | + ).append(value) |
| 185 | + |
| 186 | + return outer_fold_metrics |
| 187 | + |
| 188 | + def _score_one_inner_trial(self, model, x_folds, y_folds, metric) -> float: |
| 189 | + """The objective of an inner search: the fold loop, recording nothing. |
| 190 | +
|
| 191 | + The partitions of an inner trial belong to one outer fold. Rows written |
| 192 | + from them would sit in the same table as the rows describing the model |
| 193 | + that gets kept, indistinguishable from them and far more numerous. |
| 194 | + """ |
| 195 | + return self._score_folds(model, x_folds, y_folds, metric, record=False) |
0 commit comments