[WIP] Test time deep method

skada/deep/_test_time.py

@@ -0,0 +1,99 @@
+# Author : Maxence Barneche
+#
+# License: BSD-3-Clause
+
+import torch
+
+from skada.deep.base import DomainAwareNet
+
+
+class TestTimeCriterion(torch.nn.Module):
+    def __init__(
+        self,
+        base_criterion,
+        adapt_criterion,
+        reg=1,
+        reduction="mean",
+        train_on_target=False,
+    ):
+        super().__init__()
+        self.base_criterion = base_criterion
+        self.adapt_criterion = adapt_criterion
+        self.reg = reg
+        self.train_on_target = train_on_target
+
+        # Update the reduce parameter for both criteria if specified
+        if hasattr(self.base_criterion, "reduction"):
+            self.base_criterion.reduction = reduction
+
+        # TODO: implement losses between source and target
+        # that are sum of the losses of each sample
+        # if hasattr(self.adapt_criterion, 'reduction'):
+        #     self.adapt_criterion.reduction = reduction
+
+    def forward(
+        self,
+        y_pred,
+        y_true,
+    ):
+        """
+        Parameters
+        ----------
+        y_pred : tuple
+            This tuple comprises all the different data
+            needed to compute DA loss:
+                - y_pred : prediction of the source and target domains
+                - domain_pred : prediction of domain classifier if given
+                - features :  features of the chosen layer
+                  of source and target domains
+                - sample_domain : giving the domain of each samples
+        y_true :
+            The true labels. Available for source, masked for target.
+        """
+        y_pred, domain_pred, features, _, sample_idx = y_pred
+
+        if self.train_on_target:
+            # In finetune mode, we only compute the base loss
+            return self.base_criterion(y_pred, y_true)
+        else:
+            # In adapt mode, we compute the adaptation loss
+            return self.adapt_criterion(
+                y_s=y_true,
+                y_pred=y_pred,
+                domain_pred=domain_pred,
+                features=features,
+                sample_idx=sample_idx,
+            )
+
+
+class TestTimeNet(DomainAwareNet):
+    def __init__(self, module, criterion: "TestTimeCriterion", **kwargs):
+        super().__init__(module, criterion=criterion, **kwargs)
+
+    def fit_prepare(
+        self, X, y=None, sample_domain=None, sample_weight=None, **fit_params
+    ):
+        X = self._prepare_input(X, y, sample_domain, sample_weight)
+        X = X.select_source()
+        temp = self.criterion.train_on_target
+        self.criterion.train_on_target = False
+        self.partial_fit(X, None, **fit_params)
+        self.criterion.train_on_target = temp
+        return self
+
+    def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
+        if sample_domain is None:
+            sample_domain = -1
+        X = self._prepare_input(X, None, sample_domain, sample_weight)
+        X = X.select_target()
+        temp = self.criterion.train_on_target
+        self.criterion.train_on_target = True
+        self.partial_fit(X, None, **fit_params)
+        self.criterion.train_on_target = temp
+        return self
+
+    def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
+        X = self._prepare_input(X, y, sample_domain, sample_weight)
+        self.fit_prepare(X, **fit_params)
+        self.fit_adapt(X, **fit_params)
+        return self

skada/deep/tests/test_test_time.py

@@ -0,0 +1,41 @@
+# Author : Maxence Barneche
+#
+# License: BSD-3-Clause
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from torch.utils.data import DataLoader
+
+from skada.datasets import make_shifted_datasets
+from skada.deep._test_time import TestTimeCriterion, TestTimeNet
+from skada.deep.base import DomainAwareModule
+from skada.deep.losses import TestLoss
+from skada.deep.modules import ToyModule2D
+
+
+def test_test_time_criterion():
+    num_features = 10
+    module = ToyModule2D(num_features=num_features)
+    criterion = TestTimeCriterion(torch.nn.CrossEntropyLoss(), TestLoss())
+
+    n_samples = 20
+    dataset = make_shifted_datasets(
+        n_samples_source=n_samples,
+        n_samples_target=n_samples,
+        shift="conditional_shift",
+        noise=0.1,
+        random_state=42,
+    )
+    X, y, sample_domain = dataset
+    method = TestTimeNet(
+        DomainAwareModule(module, "dropout"),
+        iterator_train=DataLoader,
+        criterion=criterion,
+        batch_size=10,
+        max_epochs=2,
+        train_split=None,
+    )
+
+    method.fit(X, y, sample_domain)