Merge ac120e68ef into a5a162044c

library/model_utils.py

@@ -0,0 +1,18 @@
+import torch
+from torch import nn, Tensor
+
+
+class AID(nn.Module):
+    def __init__(self, p=0.9):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: Tensor):
+        if self.training:
+            pos_mask = (x >= 0) * torch.bernoulli(torch.ones_like(x) * self.p)
+            neg_mask = (x < 0) * torch.bernoulli(torch.ones_like(x) * (1 - self.p))
+            return x * (pos_mask + neg_mask)
+        else:
+            pos_part = (x >= 0) * x * self.p
+            neg_part = (x < 0) * x * (1 - self.p)
+            return pos_part + neg_part