[MRG] In DEV, reshape features to 2D instead of input (#226)

* In DEV, reshape features to 2D instead of input

* Add tests for dev scorer

* Remove raise arg in cv

* comments

* improve docstring

* use da_dataset fixture

* automatic number of channels, ...

* modify ToyCNN so that features are more than 2D

* rm not necessary reshapes

* float32

* change test name

* check features are actually more than 2D

* spaces

* rm from test_scorer_with_nd_input

---------

Co-authored-by: Antoine Collas <[email protected]>

Author: YanisLalou

skada/deep/modules.py

@@ -41,7 +41,9 @@ def forward(
 
 
 class ToyCNN(nn.Module):
-    """Toy CNN for examples and tests.
+    """Toy CNN for examples and tests on classification tasks.
+
+    Made for 2D data (e.g. time series) with shape (batch_size, n_channels, input_size).
 
     Parameters
     ----------
@@ -63,15 +65,31 @@ def __init__(
         self.feature_extractor = nn.Sequential(
             nn.Conv1d(n_channels, out_channels, kernel_size),
             nn.ReLU(),
-            nn.AvgPool1d(kernel_size),
         )
         self.num_features = self._num_features(n_channels, input_size)
-        self.fc = nn.Linear(self.num_features, n_classes)
+        self.fc = nn.Sequential(
+            nn.AdaptiveAvgPool1d(1),
+            nn.Flatten(start_dim=1),
+            nn.Linear(out_channels, n_classes),
+        )
 
     def forward(self, x, sample_weight=None):
-        """XXX add docstring here."""
+        """Forward pass of the network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, n_channels, input_size).
+        sample_weight : torch.Tensor, optional
+            Sample weights for the loss computation of shape (batch_size,).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (batch_size, n_classes).
+        """
         x = self.feature_extractor(x)
-        x = self.fc(x.flatten(start_dim=1))
+        x = self.fc(x)
         return x
 
     def _num_features(self, n_channels, input_size):

skada/deep/tests/test_deep_scorer.py

@@ -4,12 +4,13 @@
 
 import numpy as np
 import pytest
+import torch
 from sklearn.model_selection import ShuffleSplit, cross_validate
 from sklearn.preprocessing import StandardScaler
 
-from skada import make_da_pipeline
+from skada import make_da_pipeline, source_target_split
 from skada.deep import DeepCoral
-from skada.deep.modules import ToyModule2D
+from skada.deep.modules import ToyCNN, ToyModule2D
 from skada.metrics import (
     CircularValidation,
     DeepEmbeddedValidation,
@@ -87,7 +88,47 @@ def test_generic_scorer(scorer, da_dataset):
         cv=cv,
         params={"sample_domain": sample_domain},
         scoring=scorer,
-        error_score="raise",
+    )["test_score"]
+    assert scores.shape[0] == 3, "evaluate 3 splits"
+    assert np.all(~np.isnan(scores)), "all scores are computed"
+
+
+def test_dev_cnn_features_nd(da_dataset):
+    X, y, sample_domain = da_dataset.pack_train(as_sources=["s"], as_targets=["t"])
+    X = np.repeat(X[..., np.newaxis], repeats=5, axis=-1)  # Make it batched 2D data
+    X = X.astype(np.float32)
+
+    scorer = DeepEmbeddedValidation()
+    _, n_channels, input_size = X.shape
+    y_source, _ = source_target_split(y, sample_domain=sample_domain)
+    n_classes = len(np.unique(y_source))
+    module = ToyCNN(
+        n_channels=n_channels,
+        input_size=input_size,
+        n_classes=n_classes,
+        kernel_size=3,
+        out_channels=2,
+    )
+    # Assert features more than 2D
+    assert module.feature_extractor(torch.tensor(X)).ndim > 2
+
+    net = DeepCoral(
+        module,
+        reg=1,
+        layer_name="feature_extractor",
+        batch_size=10,
+        max_epochs=10,
+        train_split=None,
+    )
+
+    cv = ShuffleSplit(n_splits=3, test_size=0.3, random_state=0)
+    scores = cross_validate(
+        net,
+        X,
+        y,
+        cv=cv,
+        params={"sample_domain": sample_domain},
+        scoring=scorer,
     )["test_score"]
     assert scores.shape[0] == 3, "evaluate 3 splits"
     assert np.all(~np.isnan(scores)), "all scores are computed"

skada/metrics.py

@@ -412,7 +412,6 @@ def identity(x):
             transformer = identity
 
         X, y, sample_domain = check_X_y_domain(X, y, sample_domain, allow_nd=True)
-        X = X.reshape(X.shape[0], -1)
         source_idx = extract_source_indices(sample_domain)
         rng = check_random_state(self.random_state)
         X_train, X_val, _, y_val, _, sample_domain_val = train_test_split(

skada/tests/test_scorer.py

@@ -316,7 +316,6 @@ def test_mixval_scorer_regression(da_reg_dataset):
         ImportanceWeightedScorer(),
         PredictionEntropyScorer(),
         SoftNeighborhoodDensity(),
-        DeepEmbeddedValidation(),
         CircularValidation(),
         MixValScorer(alpha=0.55, random_state=42),
     ],