[MRG] Add SPA method

README.md

@@ -241,4 +241,6 @@ The library is distributed under the 3-Clause BSD license.
 
 [34] Jin, Ying, Wang, Ximei, Long, Mingsheng, Wang, Jianmin. [Minimum Class Confusion for Versatile Domain Adaptation](https://arxiv.org/pdf/1912.03699). ECCV, 2020. 
 
-[35] Zhang, Y., Liu, T., Long, M., & Jordan, M. I. (2019). [Bridging Theory and Algorithm for Domain Adaptation](https://arxiv.org/abs/1904.05801). In Proceedings of the 36th International Conference on Machine Learning, (pp. 7404-7413).
+[35] Zhang, Y., Liu, T., Long, M., & Jordan, M. I. (2019). [Bridging Theory and Algorithm for Domain Adaptation](https://arxiv.org/abs/1904.05801). In Proceedings of the 36th International Conference on Machine Learning, (pp. 7404-7413).
+
+[36] Xiao, Zhiqing, Wang, Haobo, Jin, Ying, Feng, Lei, Chen, Gang, Huang, Fei, Zhao, Junbo.[SPA: A Graph Spectral Alignment Perspective for Domain Adaptation](https://arxiv.org/pdf/2310.17594). In Neurips, 2023.

skada/deep/__init__.py

@@ -20,6 +20,7 @@
 from ._optimal_transport import DeepJDOT, DeepJDOTLoss
 from ._adversarial import DANN, CDAN, MDD, DANNLoss, CDANLoss, MDDLoss, ModifiedCrossEntropyLoss
 from ._class_confusion import MCC, MCCLoss
+from ._graph_alignment import SPA, SPALoss
 from ._baseline import SourceOnly, TargetOnly
 
 from . import losses
@@ -45,6 +46,8 @@
     "ModifiedCrossEntropyLoss",
     "CANLoss",
     "CAN",
+    'SPALoss',
+    'SPA',
     "SourceOnly",
     "TargetOnly",
 ]

skada/deep/_adversarial.py

@@ -53,13 +53,9 @@ def __init__(self, domain_criterion=None):
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
-        y_pred_t,
         domain_pred_s,
         domain_pred_t,
-        features_s,
-        features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         domain_label = torch.zeros(
@@ -183,13 +179,11 @@ def __init__(self, domain_criterion=None):
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
-        y_pred_t,
         domain_pred_s,
         domain_pred_t,
         features_s,
         features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         dtype = torch.float32
@@ -247,7 +241,14 @@ def __init__(
         self.max_features = max_features
         self.random_state = random_state
 
-    def forward(self, X, sample_domain=None, is_fit=False, return_features=False):
+    def forward(
+        self,
+        X,
+        sample_domain=None,
+        sample_idx=None,
+        is_fit=False,
+        return_features=False,
+    ):
         if is_fit:
             # predict
             y_pred = self.base_module_(X)
@@ -279,6 +280,7 @@ def forward(self, X, sample_domain=None, is_fit=False, return_features=False):
                 domain_pred,
                 features,
                 sample_domain,
+                sample_idx,
             )
         else:
             if return_features:
@@ -417,22 +419,20 @@ def __init__(self, gamma=4.0):
 
     def forward(
         self,
-        y_s,
         y_pred_s,
         y_pred_t,
-        disc_pred_s,
-        disc_pred_t,
-        features_s,
-        features_t,
+        domain_pred_s,
+        domain_pred_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         # TODO: handle binary classification
         # Multiclass classification
         pseudo_label_s = torch.argmax(y_pred_s, axis=-1)
         pseudo_label_t = torch.argmax(y_pred_t, axis=-1)
 
-        disc_loss_s = self.disc_criterion_s(disc_pred_s, pseudo_label_s)
-        disc_loss_t = self.disc_criterion_t(disc_pred_t, pseudo_label_t)
+        disc_loss_s = self.disc_criterion_s(domain_pred_s, pseudo_label_s)
+        disc_loss_t = self.disc_criterion_t(domain_pred_t, pseudo_label_t)
 
         # Compute the MDD loss value
         disc_loss = self.gamma * disc_loss_s - disc_loss_t

skada/deep/_baseline.py

@@ -26,13 +26,7 @@ def __init__(
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
-        y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
-        features_s,
-        features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         return 0

skada/deep/_class_confusion.py

@@ -42,13 +42,8 @@ def __init__(self, T=1, eps=1e-7):
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
         y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
-        features_s,
-        features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         loss = mcc_loss(

skada/deep/_divergence.py

@@ -46,13 +46,9 @@ def __init__(
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
-        y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
         features_s,
         features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         loss = deepcoral_loss(features_s, features_t, self.assume_centered)
@@ -132,13 +128,9 @@ def __init__(self, sigmas=None, eps=1e-7):
 
     def forward(
         self,
-        y_s,
-        y_pred_s,
-        y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
         features_s,
         features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         loss = dan_loss(features_s, features_t, sigmas=self.sigmas, eps=self.eps)
@@ -235,12 +227,9 @@ def __init__(
     def forward(
         self,
         y_s,
-        y_pred_s,
-        y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
         features_s,
         features_t,
+        **kwargs,
     ):
         loss = cdd_loss(
             y_s,

skada/deep/_graph_alignment.py

@@ -0,0 +1,201 @@
+# Author: Theo Gnassounou <[email protected]>
+#
+# License: BSD 3-Clause
+
+import torch
+
+from skada.deep.base import (
+    BaseDALoss,
+    DomainAwareCriterion,
+    DomainAwareModule,
+    DomainAwareNet,
+    DomainBalancedDataLoader,
+)
+from skada.deep.callbacks import ComputeMemoryBank, MemoryBankInit
+from skada.deep.losses import gda_loss, nap_loss
+
+from .modules import DomainClassifier
+
+
+class SPALoss(BaseDALoss):
+    """Loss SPA.
+
+    This loss tries to minimize the divergence between features with
+    adversarial method. The weights are updated to make harder
+    to classify domains (i.e., remove domain-specific features).
+
+    See [36]_ for details.
+
+    Parameters
+    ----------
+    target_criterion : torch criterion (class), default=None
+        The initialized criterion (loss) used to compute the
+        adversarial loss. If None, a BCELoss is used.
+    reg_adv : float, default=1
+        Regularization parameter for adversarial loss.
+    reg_gsa : float, default=1
+        Regularization parameter for graph alignment loss
+    reg_nap : float, default=1
+        Regularization parameter for nap loss
+
+    References
+    ----------
+    .. [36] Xiao et. al. SPA: A Graph Spectral Alignment Perspective for
+            Domain Adaptation. In Neurips, 2023.
+    """
+
+    def __init__(
+        self,
+        domain_criterion=None,
+        memory_features=None,
+        memory_outputs=None,
+        K=5,
+        reg_adv=1,
+        reg_gsa=1,
+        reg_nap=1,
+    ):
+        super().__init__()
+        if domain_criterion is None:
+            self.domain_criterion_ = torch.nn.BCELoss()
+        else:
+            self.domain_criterion_ = domain_criterion
+
+        self.reg_adv = reg_adv
+        self.reg_gsa = reg_gsa
+        self.reg_nap = reg_nap
+        self.K = K
+        self.memory_features = memory_features
+        self.memory_outputs = memory_outputs
+
+    def forward(
+        self,
+        domain_pred_s,
+        domain_pred_t,
+        features_s,
+        features_t,
+        sample_idx_t,
+        **kwargs,
+    ):
+        """Compute the domain adaptation loss"""
+        domain_label = torch.zeros(
+            (domain_pred_s.size()[0]),
+            device=domain_pred_s.device,
+        )
+        domain_label_target = torch.ones(
+            (domain_pred_t.size()[0]),
+            device=domain_pred_t.device,
+        )
+
+        # update classification function
+        loss_adv = self.domain_criterion_(
+            domain_pred_s, domain_label
+        ) + self.domain_criterion_(domain_pred_t, domain_label_target)
+
+        loss_gda = self.reg_gsa * gda_loss(features_s, features_t)
+
+        loss_pl = self.reg_nap * nap_loss(
+            features_s,
+            features_t,
+            self.memory_features,
+            self.memory_outputs,
+            K=self.K,
+            sample_idx=sample_idx_t,
+        )
+        loss = loss_adv + loss_gda + loss_pl
+        return loss
+
+
+def SPA(
+    module,
+    layer_name,
+    reg_adv=1,
+    reg_gsa=1,
+    reg_nap=1,
+    domain_classifier=None,
+    num_features=None,
+    base_criterion=None,
+    domain_criterion=None,
+    callbacks=None,
+    **kwargs,
+):
+    """Domain Adaptation with SPA.
+
+    From [36]_.
+
+    Parameters
+    ----------
+    module : torch module (class or instance)
+        A PyTorch :class:`~torch.nn.Module`. In general, the
+        uninstantiated class should be passed, although instantiated
+        modules will also work.
+    layer_name : str
+        The name of the module's layer whose outputs are
+        collected during the training.
+    reg : float, default=1
+        Regularization parameter for DA loss.
+    domain_classifier : torch module, default=None
+        A PyTorch :class:`~torch.nn.Module` used to classify the
+        domain. If None, a domain classifier is created following [1]_.
+    num_features : int, default=None
+        Size of the input of domain classifier,
+        e.g size of the last layer of
+        the feature extractor.
+        If domain_classifier is None, num_features has to be
+        provided.
+    base_criterion : torch criterion (class)
+        The base criterion used to compute the loss with source
+        labels. If None, the default is `torch.nn.CrossEntropyLoss`.
+    domain_criterion : torch criterion (class)
+        The criterion (loss) used to compute the
+        adversarial loss. If None, a BCELoss is used.
+
+    References
+    ----------
+    .. [36] Xiao et. al. SPA: A Graph Spectral Alignment Perspective for
+            Domain Adaptation. In Neurips, 2023.
+    """
+    if domain_classifier is None:
+        # raise error if num_feature is None
+        if num_features is None:
+            raise ValueError(
+                "If domain_classifier is None, num_features has to be provided"
+            )
+        domain_classifier = DomainClassifier(num_features=num_features)
+
+    if callbacks is None:
+        callbacks = [
+            ComputeMemoryBank(),
+            MemoryBankInit(),
+        ]
+    else:
+        if isinstance(callbacks, list):
+            callbacks.append(ComputeMemoryBank())
+            callbacks.append(MemoryBankInit())
+        else:
+            callbacks = [
+                callbacks,
+                ComputeMemoryBank(),
+                MemoryBankInit(),
+            ]
+    if base_criterion is None:
+        base_criterion = torch.nn.CrossEntropyLoss()
+
+    net = DomainAwareNet(
+        module=DomainAwareModule,
+        module__base_module=module,
+        module__layer_name=layer_name,
+        module__domain_classifier=domain_classifier,
+        iterator_train=DomainBalancedDataLoader,
+        criterion=DomainAwareCriterion,
+        criterion__base_criterion=base_criterion,
+        criterion__reg=1,
+        criterion__adapt_criterion=SPALoss(
+            domain_criterion=domain_criterion,
+            reg_adv=reg_adv,
+            reg_gsa=reg_gsa,
+            reg_nap=reg_nap,
+        ),
+        callbacks=callbacks,
+        **kwargs,
+    )
+    return net

skada/deep/_optimal_transport.py

@@ -51,12 +51,10 @@ def __init__(self, reg_dist=1, reg_cl=1, target_criterion=None):
     def forward(
         self,
         y_s,
-        y_pred_s,
         y_pred_t,
-        domain_pred_s,
-        domain_pred_t,
         features_s,
         features_t,
+        **kwargs,
     ):
         """Compute the domain adaptation loss"""
         loss = deepjdot_loss(