Merge pull request #9 from plutonium-239/dropout

merge dropout into main

memsave_torch/nn/Dropout.py

@@ -0,0 +1,46 @@
+"""Implementation of a memory saving Dropout (sort of).
+
+This is done by not saving the whole input/output `float32` tensor and instead just saving the `bool` mask (8bit).
+"""
+
+import torch
+import torch.nn as nn
+
+from memsave_torch.nn.functional import dropoutMemSave
+
+
+class MemSaveDropout(nn.Dropout):
+    """MemSaveDropout."""
+
+    def __init__(self, p=0.5):
+        """Inits a MemSaveDropout layer with the given params.
+
+        Args:
+            p: Probability of elements being zeroed
+        """
+        super().__init__(p)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x: Input to the network
+
+        Returns:
+            torch.Tensor: Output
+        """
+        return dropoutMemSave(x, self.p, self.train)
+
+    @classmethod
+    def from_nn_dropout(cls, dropout: nn.Dropout):
+        """Converts a nn.Dropout layer to MemSaveDropout.
+
+        Args:
+            dropout : The nn.Dropout layer
+
+        Returns:
+            obj: The MemSaveDropout object
+        """
+        obj = cls(dropout.p)
+        return obj
+

memsave_torch/nn/functional/Dropout.py

@@ -15,7 +15,8 @@ def forward(ctx, x, p, train):
         out = torch.dropout(x, p, train)
         if ctx.needs_input_grad[0]:
             ctx.p = p
-            ctx.mask = mask
+            ctx.train = train
+            ctx.rng = rng
         return out
 
     @staticmethod