From 41f3c43f3c9e9a08da4887ce2cc447e940faeaad Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:28:49 +0200 Subject: [PATCH] fix(data): reshape labels instead of unsqueeze so a column-vector y works fit() documents y as (n_samples,) or (n_samples, n_targets), but TabularDataModule.setup() applied unsqueeze(dim=1) without flattening first. An (n, 1) target therefore became a (B, 1, 1) label tensor against (B, 1) predictions: MSELoss broadcast that to (B, B, 1) and optimised an all-pairs objective, losing the per-row pairing entirely, while BCEWithLogitsLoss refused to broadcast and raised. Same data and seed, only the y shape differing, gave R2 0.996 for 1-D y versus 0.002 for y.reshape(-1, 1). The corrupted loss also drove val_loss, so early stopping and checkpoint selection were affected too. reshape(-1, 1) produces the intended (B, 1) labels for both 1-D and column-vector input. Fixes #441 Co-Authored-By: Claude Fable 5 --- deeptab/data/datamodule.py | 13 ++++--- tests/test_target_shapes.py | 69 +++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 5 deletions(-) create mode 100644 tests/test_target_shapes.py diff --git a/deeptab/data/datamodule.py b/deeptab/data/datamodule.py index f25222a..397cdd5 100644 --- a/deeptab/data/datamodule.py +++ b/deeptab/data/datamodule.py @@ -269,11 +269,14 @@ def setup(self, stage: str): if key in val_preprocessed_data: val_emb_tensors.append(torch.tensor(val_preprocessed_data[key], dtype=torch.float32)) - # Prepare labels with appropriate shape and dtype based on task + # Prepare labels with appropriate shape and dtype based on task. + # y is documented as (n_samples,) or (n_samples, n_targets); reshape + # rather than unsqueeze so an (n, 1) column vector does not become a + # 3-D label tensor that silently broadcasts against (B, 1) predictions. if self.regression: # Regression: float32, shape (batch_size, 1) - train_labels = torch.tensor(self.y_train, dtype=torch.float32).unsqueeze(dim=1) - val_labels = torch.tensor(self.y_val, dtype=torch.float32).unsqueeze(dim=1) + train_labels = torch.tensor(self.y_train, dtype=torch.float32).reshape(-1, 1) + val_labels = torch.tensor(self.y_val, dtype=torch.float32).reshape(-1, 1) else: # Classification: determine if binary or multiclass num_classes = len(np.unique(self.y_train)) # type: ignore[arg-type] @@ -283,8 +286,8 @@ def setup(self, stage: str): val_labels = torch.tensor(self.y_val, dtype=torch.long).view(-1) else: # Binary: float32, shape (batch_size, 1) - train_labels = torch.tensor(self.y_train, dtype=torch.float32).unsqueeze(dim=1) - val_labels = torch.tensor(self.y_val, dtype=torch.float32).unsqueeze(dim=1) + train_labels = torch.tensor(self.y_train, dtype=torch.float32).reshape(-1, 1) + val_labels = torch.tensor(self.y_val, dtype=torch.float32).reshape(-1, 1) self.train_dataset = TabularDataset( train_cat_tensors, diff --git a/tests/test_target_shapes.py b/tests/test_target_shapes.py new file mode 100644 index 0000000..7cee556 --- /dev/null +++ b/tests/test_target_shapes.py @@ -0,0 +1,69 @@ +"""Regression tests: a (n, 1) column-vector y must behave like a 1-D y. + +fit() documents ``y : array-like, shape (n_samples,) or (n_samples, n_targets)``, +but setup() used ``unsqueeze(dim=1)``, turning an (n, 1) target into a (B, 1, 1) +label tensor. MSELoss then broadcast against the (B, 1) predictions and optimised +an all-pairs objective (R2 collapsed from 0.996 to 0.002), while +BCEWithLogitsLoss refused to broadcast and raised. +""" + +import numpy as np +import pandas as pd +import pytest +from sklearn.metrics import r2_score + +from deeptab.configs.core import PreprocessingConfig +from deeptab.models import MLPClassifier, MLPRegressor + +QUIET = { + "accelerator": "cpu", + "devices": 1, + "enable_progress_bar": False, + "enable_model_summary": False, + "logger": False, +} + + +def _regression_data(n=80, seed=0): + rng = np.random.default_rng(seed) + X = pd.DataFrame({"a": rng.normal(size=n), "b": rng.normal(size=n)}) + y = 3.0 * X["a"].to_numpy() + 0.1 * rng.normal(size=n) + return X, y + + +@pytest.mark.parametrize("as_column_vector", [False, True]) +def test_regression_label_tensor_is_2d(as_column_vector): + X, y = _regression_data() + target = y.reshape(-1, 1) if as_column_vector else y + + model = MLPRegressor(random_state=0) + model.fit(X, target, max_epochs=1, batch_size=16, **QUIET) + model._data_module.setup("fit") # type: ignore[union-attr] + + labels = model._data_module.train_dataset.labels # type: ignore[union-attr] + assert labels.ndim == 2, f"expected (B, 1) labels, got {tuple(labels.shape)}" + + +def test_column_vector_y_trains_as_well_as_1d(): + X, y = _regression_data() + scores = {} + for tag, target in (("1d", y), ("2d", y.reshape(-1, 1))): + model = MLPRegressor( + preprocessing_config=PreprocessingConfig(numerical_preprocessing="standardization"), + random_state=0, + ) + model.fit(X, target, max_epochs=30, batch_size=16, lr=1e-2, patience=30, **QUIET) + scores[tag] = r2_score(y, model.predict(X)) + + assert scores["1d"] > 0.9, f"1-D baseline did not train: {scores}" + assert scores["2d"] == pytest.approx(scores["1d"], abs=0.1), f"column vector trained differently: {scores}" + + +def test_binary_classification_accepts_column_vector_y(): + rng = np.random.default_rng(1) + X = pd.DataFrame({"a": rng.normal(size=48), "b": rng.normal(size=48)}) + y = (X["a"].to_numpy() > 0).astype(int) + + model = MLPClassifier(random_state=0) + model.fit(X, y.reshape(-1, 1), max_epochs=1, batch_size=16, **QUIET) + assert model.predict(X).shape == (48,)