[BUG] Save/load loses estimator state, and estimators are not recognised as classifiers/regressors by sklearn

[ChrisW09]

Describe the bug

A group of related gaps in how estimator state survives load() and how estimators present themselves
to scikit-learn.

---

score() on a loaded regressor raises AttributeError: '_best_model_path' is never restored by load()

Where: deeptab/models/_mixins/predict.py (117)

_score() reads self._best_model_path unguarded, but load() builds the estimator with __new__ and restore_base_state never sets that attribute, so the very first call to Regressor.load(p).score(X, y) raises AttributeError.

Observed: AttributeError: 'MLPRegressor' object has no attribute '_best_model_path'. (predict(), evaluate(), describe(), summary(), runtime_info(), parameter_table() on the same loaded object all work, so only score() — the sklearn-contract method — is broken.)

Expected: MLPRegressor.load(path).score(X, y) returns the R² score, matching the documented score() contract; _score should tolerate a missing/None _best_model_path (the loaded weights are already the best ones).

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.mlp import MLPRegressor
from deeptab.configs import MLPConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=60), 'b': rng.normal(size=60)})
y = X['a'].values * 2
m = MLPRegressor(model_config=MLPConfig(layer_sizes=[16,8]),
                 trainer_config=TrainerConfig(max_epochs=2, batch_size=16, patience=3), random_state=42)
m.fit(X, y, accelerator='cpu', devices=1); m.save('r.deeptab')
MLPRegressor.load('r.deeptab').score(X.iloc[:20], y[:20])

---

load() nulls model_config / preprocessing_config / trainer_config / random_state, so fitting a loaded model silently retrains with library defaults

Where: deeptab/core/serialization.py (381-384)

restore_base_state hard-assigns all three config slots and random_state to None. A loaded estimator therefore forgets every training hyperparameter it was configured with, and a subsequent fit() silently falls back to defaults (batch_size 128, config lr, no seed) instead of the settings the artifact was trained under.

Observed: None None None None; after refit: batch_size 128, lr 0.0001, random_state None — versus 16 / 0.05 / 42 in the original estimator. L.get_params() also flips to the legacy flat-kwargs shape, so set_params()/clone() on a loaded estimator no longer address the configs.

Expected: The save bundle already carries config, config_kwargs, preprocessor_kwargs, optimizer_type/optimizer_kwargs and the trainer scalars; the reloaded estimator should expose an equivalent model_config/preprocessing_config/trainer_config and random_state so that fit-after-load reproduces the original training setup rather than defaults.

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.mlp import MLPRegressor
from deeptab.configs import MLPConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=60), 'b': rng.normal(size=60)})
y = X['a'].values * 2
reg = MLPRegressor(model_config=MLPConfig(layer_sizes=[16,8]),
                   trainer_config=TrainerConfig(max_epochs=2, batch_size=16, patience=3, lr=0.05, optimizer_type='SGD'),
                   random_state=42)
reg.fit(X, y, accelerator='cpu', devices=1); reg.save('r.deeptab')
L = MLPRegressor.load('r.deeptab')
print(L.trainer_config, L.preprocessing_config, L.model_config, L.random_state)
L.fit(X, y, max_epochs=3, accelerator='cpu', devices=1)
print(L._data_module.batch_size, L._task_model.lr, L.random_state)

---

Estimators declare no sklearn estimator type, so is_classifier()/is_regressor() are False — meta-estimators reject them and CV is not stratified

Where: deeptab/models/base.py (118-126)

SklearnBase inherits only BaseEstimator (no ClassifierMixin/RegressorMixin, no _estimator_type, no __sklearn_tags__ override), so sklearn.base.is_classifier(MLPClassifier()) is False. Downstream sklearn machinery that branches on estimator type then misbehaves or refuses the estimator.

Observed: False False; check_cv returns KFold (not StratifiedKFold); VotingClassifier(...).fit(...) raises ValueError: The estimator MLPClassifier should be a classifier. sklearn.utils.get_tags(clf).estimator_type is None. (Plain Pipeline, cross_val_score, cross_val_predict and GridSearchCV(model_config__…) do work.)

Expected: is_classifier(MLPClassifier()) is True and is_regressor(MLPRegressor()) is True, so classification CV is stratified by default and the estimators can be used inside VotingClassifier/other classifier-only meta-estimators — sklearn 1.9 does this via __sklearn_tags__ / the Classifier–RegressorMixin.

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from sklearn.base import is_classifier, is_regressor
from sklearn.model_selection import check_cv
from sklearn.ensemble import VotingClassifier
from sklearn.linear_model import LogisticRegression
from deeptab.models.mlp import MLPClassifier, MLPRegressor
from deeptab.configs import MLPConfig, TrainerConfig
TC = TrainerConfig(max_epochs=2, batch_size=16, patience=3)
clf = MLPClassifier(model_config=MLPConfig(layer_sizes=[8]), trainer_config=TC, random_state=0)
print(is_classifier(clf), is_regressor(MLPRegressor(model_config=MLPConfig(), trainer_config=TC)))
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=60), 'b': rng.normal(size=60)}); y = (X['a']>0).astype(int).values
print(type(check_cv(3, y, classifier=is_classifier(clf))).__name__)
VotingClassifier([('a', clf), ('b', LogisticRegression())], voting='soft').fit(X, y)

---

observability_config is an init parameter but is missing from get_params()/set_params(), so clone() silently drops all observability

Where: deeptab/models/base.py (258-280)

The overridden get_params returns only model_config/preprocessing_config/trainer_config/random_state. Since clone() reconstructs from get_params(), every clone (GridSearchCV, cross_val_score, Pipeline, clone()) loses the ObservabilityConfig — no run directory, no lifecycle log, no summary.json, no MLflow logging — with no warning. set_params(observability_config=...) is likewise a silent no-op.

Observed: get_params keys: ['model_config', 'preprocessing_config', 'random_state', 'trainer_config'] — no 'observability_config'. After fit: original run_dir = 'obsclone2/runs/e/20260727_220657_422ce63a', clone run_dir = None. c3.set_params(observability_config=obs) leaves _observability_config unset.

Expected: sklearn's contract is that get_params() returns every __init__ parameter, so clone(est) reproduces the estimator's configuration — observability should survive cloning (or at minimum set_params(observability_config=...) should apply it).

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from sklearn.base import clone
from deeptab.core.observability import ObservabilityConfig
from deeptab.models.mlp import MLPClassifier
from deeptab.configs import MLPConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=48), 'b': rng.normal(size=48)}); y = (X['a']>0).astype(int).values
obs = ObservabilityConfig(root_dir='obsclone2', experiment_name='e')   # structured_logging=False, no optional deps
c = MLPClassifier(model_config=MLPConfig(layer_sizes=[8]),
                  trainer_config=TrainerConfig(max_epochs=1, batch_size=16, patience=3),
                  observability_config=obs, random_state=0)
print(sorted(c.get_params(deep=False)))
c2 = clone(c)
c.fit(X, y, accelerator='cpu', devices=1); c2.fit(X, y, accelerator='cpu', devices=1)
print(c._run_dir, '|', c2._run_dir)

---

dataloader_kwargs are forwarded to the validation dataloader, so drop_last=True silently truncates validation and can crash fit()

Where: deeptab/data/datamodule.py (413-421)

dataloader_kwargs from fit() are stored on the datamodule and splatted into val_dataloader() (and predict_dataloader()/test_dataloader()) as well as train_dataloader(). With the common drop_last=True the validation loader drops its trailing partial batch; when the validation split is smaller than one batch it yields zero batches, val_loss is never logged, and EarlyStopping aborts the run with a message that points nowhere near the cause.

Observed: RuntimeError: Early stopping conditioned on metric val_loss which is not available. Pass in or modify your EarlyStopping callback to use any of the following: train_loss, train_loss_step, train_loss_epoch. With n=100 (20 val rows) fit succeeds but val_dataloader() yields 1 batch covering only 16 of the 20 validation rows, i.e. val_loss/checkpoint selection silently use 80% of the validation set.

Expected: drop_last is a training-loop concern; it should apply to the train dataloader only, so validation/prediction always see every row and fit() does not fail with an unrelated EarlyStopping error.

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.mlp import MLPRegressor
from deeptab.configs import MLPConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=64), 'b': rng.normal(size=64)})
y = X['a'].values * 2
m = MLPRegressor(model_config=MLPConfig(layer_sizes=[16]),
                 trainer_config=TrainerConfig(max_epochs=1, batch_size=16, patience=3), random_state=0)
m.fit(X, y, dataloader_kwargs={'drop_last': True}, accelerator='cpu', devices=1)   # 13 val rows < batch 16

---

Expected behavior
load() should restore enough state that the documented public methods work (score() in particular),
and the estimators should declare their sklearn estimator type so meta-estimators and stratified CV
behave. Note #410 covers the separate get_params hole for default-constructed estimators; the
observability_config omission here affects the split-config path too.

Screenshots
n/a

Desktop (please complete the following information):

• OS: macOS (Darwin 25.5.0, arm64)
• Python version: 3.11.15
• deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.