Update with new Function API

Author: chsasank

advanced_source/numpy_extensions_tutorial.py

@@ -86,17 +86,21 @@ def incorrect_fft(input):
 
 
 class ScipyConv2dFunction(Function):
-
-    def forward(self, input, filter):
+    @staticmethod
+    def forward(ctx, input, filter):
         result = correlate2d(input.numpy(), filter.numpy(), mode='valid')
-        self.save_for_backward(input, filter)
+        ctx.save_for_backward(input, filter)
         return torch.FloatTensor(result)
 
-    def backward(self, grad_output):
-        input, filter = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, filter = ctx.saved_tensors
+        grad_output = grad_output.data
         grad_input = convolve2d(grad_output.numpy(), filter.t().numpy(), mode='full')
         grad_filter = convolve2d(input.numpy(), grad_output.numpy(), mode='valid')
-        return torch.FloatTensor(grad_input), torch.FloatTensor(grad_filter)
+
+        return Variable(torch.FloatTensor(grad_input)), \
+            Variable(torch.FloatTensor(grad_filter))
 
 
 class ScipyConv2d(Module):
@@ -106,7 +110,7 @@ def __init__(self, kh, kw):
         self.filter = Parameter(torch.randn(kh, kw))
 
     def forward(self, input):
-        return ScipyConv2dFunction()(input, self.filter)
+        return ScipyConv2dFunction.apply(input, self.filter)
 
 ###############################################################
 # **Example usage:**

beginner_source/examples_autograd/two_layer_net_custom_function.py

@@ -23,22 +23,25 @@ class MyReLU(torch.autograd.Function):
     which operate on Tensors.
     """
 
-    def forward(self, input):
+    @staticmethod
+    def forward(ctx, input):
         """
-        In the forward pass we receive a Tensor containing the input and return a
-        Tensor containing the output. You can cache arbitrary Tensors for use in the
-        backward pass using the save_for_backward method.
+        In the forward pass we receive a Tensor containing the input and return
+        a Tensor containing the output. ctx is a context object that can be used
+        to stash information for backward computation. You can cache arbitrary
+        objects for use in the backward pass using the ctx.save_for_backward method.
         """
-        self.save_for_backward(input)
+        ctx.save_for_backward(input)
         return input.clamp(min=0)
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(ctx, grad_output):
         """
         In the backward pass we receive a Tensor containing the gradient of the loss
         with respect to the output, and we need to compute the gradient of the loss
         with respect to the input.
         """
-        input, = self.saved_tensors
+        input, = ctx.saved_tensors
         grad_input = grad_output.clone()
         grad_input[input < 0] = 0
         return grad_input
@@ -61,8 +64,8 @@ def backward(self, grad_output):
 
 learning_rate = 1e-6
 for t in range(500):
-    # Construct an instance of our MyReLU class to use in our network
-    relu = MyReLU()
+    # To apply our Function, we use Function.apply method. We alias this as 'relu'.
+    relu = MyReLU.apply
 
     # Forward pass: compute predicted y using operations on Variables; we compute
     # ReLU using our custom autograd operation.