Upload 3 files

gradient_reversal/__init__.py

@@ -0,0 +1,2 @@
+from .functional import revgrad
+from .module import GradientReversal

gradient_reversal/functional.py

@@ -0,0 +1,16 @@
+from torch.autograd import Function
+
+class GradientReversal(Function):
+    @staticmethod
+    def forward(ctx, x, alpha):
+        ctx.save_for_backward(x, alpha)
+        return x
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = None
+        _, alpha = ctx.saved_tensors
+        if ctx.needs_input_grad[0]:
+            grad_input = - alpha*grad_output
+        return grad_input, None
+revgrad = GradientReversal.apply

gradient_reversal/module.py

@@ -0,0 +1,11 @@
+from .functional import revgrad
+import torch
+from torch import nn
+
+class GradientReversal(nn.Module):
+    def __init__(self, alpha):
+        super().__init__()
+        self.alpha = torch.tensor(alpha, requires_grad=False)
+
+    def forward(self, x):
+        return revgrad(x, self.alpha)