From a86b4c123c69ae798e9ea62e153bae9e7667265f Mon Sep 17 00:00:00 2001 From: Maarmapa Date: Thu, 10 Sep 2026 18:33:02 +0000 Subject: [PATCH] fix(models): read MLPRegression hyperparameters from instance attributes MLPRegression kept its configuration in two places: `self.params`, the dict of construction kwargs, which is what `train`, `save` and `load` read; and instance attributes, which is what ModelFactory._process_param and both optimizers write via `setattr(obj, key, value)`. A search over hidden_size, activation, learning_rate or epochs therefore never reached the network: every trial rebuilt the construction-time model, the study reported a best trial, and the model it produced had never been trained with the values that won. Nothing failed while that was true. Mirror the configuration onto instance attributes in `__init__` and read those everywhere. Two smaller defects fall out of the same split and are fixed here too: - `hidden_size` fell back to 100 in `train` and to 5 in `load`, against a schema that declares 16. `_CONFIG_DEFAULTS` is now the single home for every fallback, and a test pins it to the schema field by field. - `save` persisted the construction kwargs next to post-search weights, so reloading rebuilt a network of the wrong width. It now records the configuration the model actually ran with. Also fixes an AttributeError that made the model unusable outside ModelFactory: `train` reads `self.log_train_every_n_steps` and its three siblings, which `__init__` never set. Adds tests/back/models/test_mlp_regression.py, including one that drives a real OptunaOptimizer over hidden_size and asserts the widths of the networks actually built equal the widths Optuna suggested. Closes #842 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BzZ6td4VJdRDe76UndWZY3 --- .../models/scikit_learn/mlp_regression.py | 71 +++++- tests/back/models/test_mlp_regression.py | 235 ++++++++++++++++++ 2 files changed, 296 insertions(+), 10 deletions(-) create mode 100644 tests/back/models/test_mlp_regression.py diff --git a/DashAI/back/models/scikit_learn/mlp_regression.py b/DashAI/back/models/scikit_learn/mlp_regression.py index 8cba2c313..90ae05e3f 100644 --- a/DashAI/back/models/scikit_learn/mlp_regression.py +++ b/DashAI/back/models/scikit_learn/mlp_regression.py @@ -324,6 +324,25 @@ class MLPRegression(CategoricalEncoderMixin, RegressionModel): COLOR: str = "#FF7043" ICON: str = "Psychology" + #: Fallback for every configurable field, in one place so ``__init__``, + #: ``train``, ``save`` and ``load`` cannot disagree about what an unset + #: value means -- ``hidden_size`` used to default to 100 when training and + #: to 5 when reloading, so a checkpoint written without that key came back + #: as a different network. Each entry is the value declared in + #: ``MLPRegressorSchema`` (``fixed`` for a search space, ``placeholder`` + #: otherwise), and a test asserts that it stays that way. + _CONFIG_DEFAULTS = { + "hidden_size": 16, + "activation": "relu", + "learning_rate": 0.001, + "epochs": 20, + "batch_size": 32, + "log_train_every_n_epochs": 1, + "log_train_every_n_steps": None, + "log_validation_every_n_epochs": 1, + "log_validation_every_n_steps": None, + } + def __init__(self, **kwargs) -> None: """Initialize the MLP regressor and set up the inner PyTorch module class. @@ -389,6 +408,20 @@ def forward(self, x): self.mlp = MLP self.params = kwargs + + # Mirror the configuration onto instance attributes, and read those -- + # never ``self.params`` -- everywhere the model is used. + # + # The optimizers assign each searched hyperparameter with + # ``setattr(model, key, value)`` once per trial (see + # ``OptunaOptimizer.optimize`` and ``HyperOptOptimizer.optimize``), and + # ``ModelFactory`` does the same for the fixed ones. A ``train`` that + # read ``self.params`` therefore trained the construction-time values on + # every trial: the search ran, the study reported a best trial, and the + # network that came out of it had never seen the values that won. + for name, default in self._CONFIG_DEFAULTS.items(): + setattr(self, name, kwargs.get(name, default)) + self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if DEVICE_TO_IDX.get(kwargs.get("device"), -1) >= 0 @@ -440,18 +473,16 @@ def train( # 2. Init Model & Optimizer self.model = self.mlp( input_dim=X_tensor.shape[1], - hidden_size=self.params.get("hidden_size", 100), - activation_name=self.params.get("activation", "relu"), + hidden_size=self.hidden_size, + activation_name=self.activation, ).to(self.device) - optimizer = torch.optim.Adam( - self.model.parameters(), lr=self.params.get("learning_rate", 0.001) - ) + optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate) criterion = torch.nn.MSELoss() # 3. Training Loop using Epochs - total_epochs = self.params.get("epochs", 3) - batch_size = self.params.get("batch_size") + total_epochs = self.epochs + batch_size = self.batch_size if batch_size is None or batch_size > X_tensor.size(0): batch_size = X_tensor.size(0) @@ -572,6 +603,26 @@ def predict_prepared(self, features) -> "ndarray": with torch.no_grad(): return self.model(x_tensor).cpu().numpy().flatten() + def _current_params(self) -> dict: + """Return the configuration the model is actually running with. + + Read from the instance attributes rather than from the construction + kwargs, so a checkpoint taken after hyperparameter optimization records + the values the optimizer chose. Saving ``self.params`` instead wrote + the pre-search configuration next to post-search weights, and reloading + that checkpoint rebuilt a network of the wrong width -- which surfaces + as a shape mismatch in ``load_state_dict``, far from its cause. + + Returns + ------- + dict + The construction kwargs, with every field of + ``_CONFIG_DEFAULTS`` overwritten by its current value. + """ + params = dict(self.params) + params.update({name: getattr(self, name) for name in self._CONFIG_DEFAULTS}) + return params + def save(self, filename: str) -> None: """Save the trained model weights and configuration to disk. @@ -585,7 +636,7 @@ def save(self, filename: str) -> None: torch.save( { "state": self.model.state_dict(), - "params": self.params, + "params": self._current_params(), "input_dim": self.model.model[0].in_features, "encodings": self.encodings, "one_hot_encoder": self.one_hot_encoder, @@ -619,8 +670,8 @@ def load(filename: str) -> "MLPRegression": # Rebuild the model architecture using saved input_dim instance.model = instance.mlp( input_dim=data["input_dim"], - hidden_size=instance.params.get("hidden_size", 5), - activation_name=instance.params.get("activation", "relu"), + hidden_size=instance.hidden_size, + activation_name=instance.activation, ).to(instance.device) # Load the trained weights diff --git a/tests/back/models/test_mlp_regression.py b/tests/back/models/test_mlp_regression.py new file mode 100644 index 000000000..8a1c42a0b --- /dev/null +++ b/tests/back/models/test_mlp_regression.py @@ -0,0 +1,235 @@ +"""The hyperparameters an optimizer suggests must reach the network built. + +``MLPRegression`` kept its configuration twice: in ``self.params``, the dict of +construction kwargs, and as instance attributes, which is where both +``ModelFactory`` and the optimizers write. ``train`` read the dict, so every +trial of a search rebuilt the same network out of the construction-time values. +The search ran, the study reported a best trial, and the model that came out of +it had never been trained with the values that won. + +Nothing failed while that was true, which is why it needed a test rather than a +type: a study with all-identical trials is indistinguishable from a study on a +flat objective unless you look at what was built. + +The dataset is 120 rows of three features and the networks are a handful of +epochs wide on purpose. This asserts wiring, not accuracy, and it should not +cost CI a minute. +""" + +import numpy as np +import pyarrow as pa +import pytest + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.evaluation.holdout import HoldoutEvaluationStrategy +from DashAI.back.metrics.regression.mse import MSE +from DashAI.back.models.scikit_learn.mlp_regression import MLPRegression +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.types.value_types import Float + +EPOCHS = 3 +N_TRIALS = 4 + +#: The width is searched over an interval wide enough that four independent +#: draws landing on the same value is not a plausible accident. +HIDDEN_SIZE_SPACE = (2, 40) + + +def _holdout_evaluate(model, input_dataset, output_dataset, metric): + """The real holdout evaluation path, on a strategy with no factory. + + ``evaluate`` reads which partitions its strategy scores, so it needs a real + instance rather than None for self. Building one through ``__init__`` would + need a ``ModelFactory`` this test does not have, and does not need: the + only thing read off the instance is a class attribute. + """ + strategy = HoldoutEvaluationStrategy.__new__(HoldoutEvaluationStrategy) + return strategy.evaluate(model, input_dataset, output_dataset, metric) + + +@pytest.fixture(scope="module", name="regression_splits") +def fixture_regression_splits(): + """A small synthetic regression dataset, split train/validation/test.""" + import pandas as pd + + rng = np.random.default_rng(0) + n = 120 + frame = pd.DataFrame( + { + "x0": rng.uniform(-2, 2, n), + "x1": rng.uniform(-2, 2, n), + "x2": rng.uniform(-2, 2, n), + } + ).astype("float64") + target = pd.DataFrame( + {"target": 3.0 * frame["x0"] - 2.0 * frame["x1"] + rng.normal(0, 0.2, n)} + ).astype("float64") + + feature_types = {c: Float(arrow_type=pa.float64()) for c in frame.columns} + target_types = {"target": Float(arrow_type=pa.float64())} + + def part(start, stop): + return ( + to_dashai_dataset(frame.iloc[start:stop], types=feature_types), + to_dashai_dataset(target.iloc[start:stop], types=target_types), + ) + + x_train, y_train = part(0, 80) + x_validation, y_validation = part(80, 100) + x_test, y_test = part(100, 120) + + return ( + {"train": x_train, "validation": x_validation, "test": x_test}, + {"train": y_train, "validation": y_validation, "test": y_test}, + ) + + +def _model(**overrides): + """An ``MLPRegression`` with the runtime state a strategy expects.""" + params = {"epochs": EPOCHS, "hidden_size": 8, "learning_rate": 0.01} + params.update(overrides) + model = MLPRegression(**params) + model.run_id = 1 + model.x_data = None + model.y_data = None + model.train_metrics = None + model.validation_metrics = None + model.test_metrics = None + return model + + +def _record_widths(model): + """Record the hidden width of every network ``train`` actually builds. + + Reading ``model.hidden_size`` afterwards would prove nothing: the optimizer + writes that attribute itself, so it holds the last suggested value whether + or not ``train`` ever read it. What has to be observed is the network. + """ + widths = [] + original = model.mlp + + def spy(input_dim, hidden_size, activation_name): + widths.append(hidden_size) + return original(input_dim, hidden_size, activation_name) + + model.mlp = spy + return widths + + +def test_suggested_values_reach_training(regression_splits): + """Every trial must train the width Optuna suggested for it. + + This is the assertion that fails on the unfixed model: the widths built are + ``[8, 8, 8, 8]`` -- the construction-time value -- while the study reports + four different suggestions. + """ + x, y = regression_splits + model = _model() + widths = _record_widths(model) + + optimizer = OptunaOptimizer( + n_trials=N_TRIALS, sampler="RandomSampler", pruner="None" + ) + optimizer.optimize( + model, + x, + y, + [(model, "hidden_size", HIDDEN_SIZE_SPACE, "integer")], + {"class": MSE, "metadata": {"maximize": False}}, + _holdout_evaluate, + ) + + suggested = [trial.params["hidden_size"] for trial in optimizer.study.trials] + assert widths == suggested, ( + f"the networks trained were {widths} wide, but Optuna suggested " + f"{suggested}. A width the search never chose means the trial scored a " + "model built from the construction-time configuration, so the study's " + "best trial describes a model that was never trained." + ) + + +def test_a_written_attribute_reaches_training(regression_splits): + """The same guarantee without Optuna, as the cheap negative control. + + ``ModelFactory`` and both optimizers configure a model the same way: plain + ``setattr`` on the instance. If this ever passes while the test above fails, + the break is in the optimizer, not in the model. + """ + x, y = regression_splits + model = _model(hidden_size=8) + widths = _record_widths(model) + + model.hidden_size = 23 + model.train(x["train"], y["train"]) + + assert widths == [23], ( + f"training built a network {widths} wide after the width was set to 23" + ) + + +def test_defaults_match_the_declared_schema(): + """The fallbacks must be the values the schema declares, all of them. + + ``hidden_size`` used to fall back to 100 when training and to 5 when + reloading, against a schema that declares 16: three answers to one + question, and a checkpoint saved without that key came back as a different + network. Pinning the table to the schema is what stops them drifting apart + again, and it also fails when a field is added to the schema and forgotten + here. + """ + properties = MLPRegression.get_schema()["properties"] + + # ``device`` is resolved into a torch device string by ``__init__`` rather + # than mirrored, so it is the one field that is not in the table. + assert set(MLPRegression._CONFIG_DEFAULTS) == set(properties) - {"device"}, ( + "the defaults table and the schema disagree about which fields exist" + ) + + for name, default in MLPRegression._CONFIG_DEFAULTS.items(): + placeholder = properties[name]["placeholder"] + # A search space declares its default inside the optimize envelope. + declared = ( + placeholder["fixed_value"] + if isinstance(placeholder, dict) and "fixed_value" in placeholder + else placeholder + ) + assert default == declared, ( + f"'{name}' falls back to {default!r} but the schema declares {declared!r}" + ) + + +def test_the_schema_defaults_are_what_an_unconfigured_model_uses(): + """And the table is what the attributes are actually built from.""" + model = MLPRegression() + for name, default in MLPRegression._CONFIG_DEFAULTS.items(): + assert getattr(model, name) == default, f"'{name}' was not mirrored" + + +def test_checkpoint_records_the_configuration_that_was_trained( + regression_splits, tmp_path +): + """A checkpoint must describe the network it holds the weights of. + + ``save`` wrote the construction kwargs, so a model saved after a search + stored the pre-search width beside post-search weights. Reloading it + rebuilt a network of the wrong width, which surfaces as a shape mismatch + inside ``load_state_dict`` -- an error that names a tensor, not the reason. + """ + x, y = regression_splits + model = _model(hidden_size=8) + + # What the optimizer does to a model between construction and training. + model.hidden_size = 19 + model.train(x["train"], y["train"]) + + path = tmp_path / "mlp.pt" + model.save(str(path)) + restored = MLPRegression.load(str(path)) + + assert restored.hidden_size == 19 + assert restored.model.model[0].out_features == 19 + + original = model.predict(x["test"]) + assert np.allclose(original, restored.predict(x["test"]), atol=1e-6), ( + "the reloaded model does not reproduce the predictions of the saved one" + )