init

daln/grl.py

@@ -0,0 +1,77 @@
+from typing import Optional, Any, Tuple
+import numpy as np
+import torch.nn as nn
+from torch.autograd import Function
+import torch
+
+
+class GradientReverseFunction(Function):
+
+    @staticmethod
+    def forward(ctx: Any, input: torch.Tensor, coeff: Optional[float] = 1.) -> torch.Tensor:
+        ctx.coeff = coeff
+        output = input * 1.0
+        return output
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, Any]:
+        return grad_output.neg() * ctx.coeff, None
+
+
+class GradientReverseLayer(nn.Module):
+    def __init__(self):
+        super(GradientReverseLayer, self).__init__()
+
+    def forward(self, *input):
+        return GradientReverseFunction.apply(*input)
+
+
+class WarmStartGradientReverseLayer(nn.Module):
+    """Gradient Reverse Layer :math:`\mathcal{R}(x)` with warm start
+
+        The forward and backward behaviours are:
+
+        .. math::
+            \mathcal{R}(x) = x,
+
+            \dfrac{ d\mathcal{R}} {dx} = - \lambda I.
+
+        :math:`\lambda` is initiated at :math:`lo` and is gradually changed to :math:`hi` using the following schedule:
+
+        .. math::
+            \lambda = \dfrac{2(hi-lo)}{1+\exp(- α \dfrac{i}{N})} - (hi-lo) + lo
+
+        where :math:`i` is the iteration step.
+
+        Args:
+            alpha (float, optional): :math:`α`. Default: 1.0
+            lo (float, optional): Initial value of :math:`\lambda`. Default: 0.0
+            hi (float, optional): Final value of :math:`\lambda`. Default: 1.0
+            max_iters (int, optional): :math:`N`. Default: 1000
+            auto_step (bool, optional): If True, increase :math:`i` each time `forward` is called.
+              Otherwise use function `step` to increase :math:`i`. Default: False
+        """
+
+    def __init__(self, alpha: Optional[float] = 1.0, lo: Optional[float] = 0.0, hi: Optional[float] = 1.,
+                 max_iters: Optional[int] = 1000., auto_step: Optional[bool] = False):
+        super(WarmStartGradientReverseLayer, self).__init__()
+        self.alpha = alpha
+        self.lo = lo
+        self.hi = hi
+        self.iter_num = 0
+        self.max_iters = max_iters
+        self.auto_step = auto_step
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        """"""
+        coeff =float(
+            2.0 * (self.hi - self.lo) / (1.0 + np.exp(-self.alpha * self.iter_num / self.max_iters))
+            - (self.hi - self.lo) + self.lo
+        )
+        if self.auto_step:
+            self.step()
+        return GradientReverseFunction.apply(input, coeff)
+
+    def step(self):
+        """Increase iteration number :math:`i` by 1"""
+        self.iter_num += 1

daln/nwd.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from daln.grl import WarmStartGradientReverseLayer
+
+
+class NuclearWassersteinDiscrepancy(nn.Module):
+    def __init__(self, classifier: nn.Module):
+        super(NuclearWassersteinDiscrepancy, self).__init__()
+        self.grl = WarmStartGradientReverseLayer(alpha=1., lo=0., hi=1., max_iters=1000, auto_step=True)
+        self.classifier = classifier
+
+    @staticmethod
+    def n_discrepancy(y_s: torch.Tensor, y_t: torch.Tensor) -> torch.Tensor:
+        pre_s, pre_t = F.softmax(y_s, dim=1), F.softmax(y_t, dim=1)
+        loss = (-torch.norm(pre_t, 'nuc') + torch.norm(pre_s, 'nuc')) / y_t.shape[0]
+        return loss
+
+    def forward(self, f: torch.Tensor) -> torch.Tensor:
+        f_grl = self.grl(f)
+        y = self.classifier(f_grl)
+        y_s, y_t = y.chunk(2, dim=0)
+
+        loss = self.n_discrepancy(y_s, y_t)
+        return loss