Hi all,
I am trying to write a user function using Cupy to interface with Pytorch using the tutorial shown here Numpy Extensions Tutorial.
I have checked that the output of the convolution is working. However, the output of the gradient is not calculated. I assume I have overlooked something.
My code is here:
from numpy import flip
import numpy as np
import cupy as cp
import sigpy
import operator
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter
from torch.utils.dlpack import to_dlpack
from torch.utils.dlpack import from_dlpack
from torch.autograd import Function
class ExampleFunction(Function):
@staticmethod
def forward(ctx, input, filter, bias):
# detach so we can cast to NumPy
ctx.save_for_backward(input, filter, bias)
input, filter, bias = sigpy.from_pytorch(input), sigpy.from_pytorch(filter), sigpy.from_pytorch(bias)
result = sigpy.convolve(input, flip(flip(filter, axis=0), axis=1), mode='valid')
result = cp.add(result, bias)
result=sigpy.to_pytorch(result, requires_grad=True)
return result
@staticmethod
def backward(ctx, grad_output):
input, filter, bias = ctx.saved_tensors
grad_output, input, filter, bias = sigpy.from_pytorch(grad_output), sigpy.from_pytorch(input), sigpy.from_pytorch(filter), sigpy.from_pytorch(bias)
grad_bias = cp.sum(grad_output, keepdims=True)
grad_input = sigpy.convolve(grad_output, filter, mode='full')
grad_filter = sigpy.convolve(input, flip(flip(grad_output, axis=0), axis=1), mode='valid')
grad_input = sigpy.to_pytorch(grad_input, requires_grad=True)
grad_bias = sigpy.to_pytorch(grad_bias, requires_grad=True)
grad_filter = sigpy.to_pytorch(grad_filter, requires_grad=True)
return grad_input, grad_filter.to(torch.float), grad_bias.to(torch.float)
class Example(Module):
def __init__(self, filter_width, filter_height):
super(Example, self).__init__()
self.filter = Parameter(torch.randn(filter_width, filter_height))
self.bias = Parameter(torch.randn(1, 1))
def forward(self, input):
return ExampleFunction.apply(input, self.filter, self.bias)
torch.cuda.is_available()
if torch.cuda.is_available():
dev = "cuda:0"
else:
dev = "cpu"
module = Example(3, 3)
module = module.to('cuda:0')
print("Filter and bias: ", list(module.parameters()))
input = torch.randn(10, 10, requires_grad=True).cuda()
output = module(input)
print("Output from the convolution2: ", output)
rand = torch.randn(8, 8).cuda()
output.backward(rand)
print("Gradient for the input map: ", input.grad)
This outputs the following:
Filter and bias: [Parameter containing:
tensor([[ 0.5030, -0.7362, -0.3410],
[ 0.2329, 0.3117, -0.3578],
[ 0.2351, 0.2211, 0.8780]], device=‘cuda:0’, requires_grad=True), Parameter containing:
tensor([[0.6562]], device=‘cuda:0’, requires_grad=True)]
Output from the convolution2: tensor([[ 0.1825, 2.5443, -0.1493, 1.2438, 2.3477, 3.8454, 1.6594, -0.8285],
[-0.0698, 1.2464, 2.1476, -1.8530, 0.5620, 2.8346, 2.4183, 0.1029],
[ 1.9444, 0.4215, 1.8121, -0.5721, -1.5620, -0.5422, 0.8763, 2.5569],
[-0.0357, 2.7819, 0.8434, 2.2674, 1.8015, 0.2055, -1.6885, 1.8099],
[ 0.9745, -0.5870, 1.4823, 2.2904, 3.4809, -1.1498, 1.9393, 0.4010],
[ 1.2713, 1.5730, 0.6142, 0.4776, 0.6151, 0.0946, -0.1609, 3.4696],
[-1.2004, -0.1432, 0.6048, 0.1454, -0.0348, 1.5170, 3.0961, -1.0933],
[ 1.4833, -1.0142, -0.7154, 0.7169, -2.2208, 1.5966, 0.0809, -1.2016]],
device=‘cuda:0’, grad_fn=)
Gradient for the input map: None
Can someone point out why this is not working?
Cheers!