Add SkorchSupervisedScorer to move tensors to cpu before scoring

Author: mdbenito

docs/value/index.md

@@ -127,8 +127,8 @@ tensors for data valuation. The implementation follows these key principles:
 1. **Type Preservation**: The valuation methods maintain the input data type
    throughout computations, whether you provide NumPy arrays or PyTorch tensors
    when constructing the [Dataset][pydvl.valuation.dataset.Dataset].
-2. **Transparent Usage**: The API remains the same regardless of the input type.
-   Simply provide your data as tensors. The main difference is that the torch
+2. **Transparent Usage**: The API remains the same regardless of the input type -
+   simply provide your data as tensors. The main difference is that the torch
    model must be wrapped in a class compatible with the protocol
    [TorchSupervisedModel][pydvl.valuation.types.TorchSupervisedModel].
      !!! tip "Wrapping torch models"

notebooks/support/banzhaf.py

@@ -10,7 +10,6 @@
 
 from pydvl.utils import timed
 from pydvl.utils.monitor import end_memory_monitoring, start_memory_monitoring
-from pydvl.valuation.types import TorchSupervisedModel
 
 from .datasets import load_digits_dataset
 
@@ -70,7 +69,7 @@ def forward(self, x):
         return self.layers(x)
 
 
-class TorchClassifierModel(TorchSupervisedModel):
+class TorchClassifierModel:
     """This class wraps a torch classification model to comply with the
     [SupervisedModel][pydvl.utils.types.SupervisedModel] interface expected by pyDVL,
     and takes care of the training and evaluation of the model.
@@ -286,7 +285,7 @@ def run(_config):
         verbose=False,
     )
 
-    # scorer = SupervisedScorer(model, test, default=0.0, range=(0.0, 1.0))
+    # scorer = SkorchSupervisedScorer(model, test, default=0.0, range=(0.0, 1.0))
 
     # utility = ModelUtility(
     #     model,

src/pydvl/valuation/scorers/__init__.py

@@ -14,4 +14,5 @@
 from .base import *
 from .classwise import *
 from .supervised import *
+from .torchscorer import *
 from .utils import *

src/pydvl/valuation/scorers/torchscorer.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from pydvl.valuation.scorers.supervised import SupervisedScorer
+from pydvl.valuation.types import SkorchSupervisedModel
+
+__all__ = ["SkorchSupervisedScorer"]
+
+logger = logging.getLogger(__name__)
+
+
+class SkorchSupervisedScorer(SupervisedScorer[SkorchSupervisedModel, torch.Tensor]):
+    """Scorer for Skorch models.
+
+    Because skorch models scorer() requires a numpy array to test against, this
+    class moves tensors to cpu before scoring.
+    """
+
+    def __call__(self, model: SkorchSupervisedModel) -> float:
+        x, y = self.test_data.data()
+        if torch.is_tensor(y):
+            y = y.cpu().numpy()
+        return float(self._scorer(model, x, y))

src/pydvl/valuation/types.py

@@ -56,7 +56,7 @@
     "SampleT",
     "SemivalueCoefficient",
     "SupervisedModel",
-    "TorchSupervisedModel",
+    "SkorchSupervisedModel",
     "UtilityEvaluation",
     "ValueUpdate",
     "ValueUpdateT",
@@ -304,7 +304,7 @@ class SupervisedModel(Protocol[ArrayT, ArrayRetT]):
     `score()`.
     """
 
-    def fit(self, x: ArrayT, y: ArrayT | None):
+    def fit(self, x: ArrayT, y: ArrayT):
         """Fit the model to the data
 
         Args:
@@ -324,7 +324,7 @@ def predict(self, x: ArrayT) -> ArrayRetT:
         """
         pass
 
-    def score(self, x: ArrayT, y: ArrayT | None) -> float:
+    def score(self, x: ArrayT, y: ArrayT) -> float:
         """Compute the score of the model given test data
 
         Args:
@@ -370,15 +370,15 @@ def predict(self, x: ArrayT) -> ArrayRetT:
 
 
 @runtime_checkable
-class TorchSupervisedModel(Protocol):
+class SkorchSupervisedModel(Protocol[ArrayT]):
     """This is the standard sklearn Protocol with the methods `fit()`, `predict()`
     and `score()`, but accepting Tensors and with any additional info required.
     It is compatible with [skorch.net.NeuralNet][].
     """
 
     device: str | torch_mod.device
 
-    def fit(self, x: Tensor, y: Tensor | None):
+    def fit(self, x: ArrayT, y: Tensor):
         """Fit the model to the data
 
         Args:
@@ -387,7 +387,7 @@ def fit(self, x: Tensor, y: Tensor | None):
         """
         ...
 
-    def predict(self, x: Tensor) -> Tensor:
+    def predict(self, x: ArrayT) -> NDArray:
         """Compute predictions for the input
 
         Args:
@@ -398,7 +398,7 @@ def predict(self, x: Tensor) -> Tensor:
         """
         ...
 
-    def score(self, x: Tensor, y: Tensor | None) -> float:
+    def score(self, x: ArrayT, y: NDArray) -> float:
         """Compute the score of the model given test data
 
         Args: