Hello everyone,
I’m trying to implement a custom convolution layer. However, I only want to custom the forward function, and leave the backward function unchanged.
Following the guidance from https://pytorch.org/docs/stable/notes/extending.html, I think the implementation would be something like
class MyConv2d(Function):
def forward(ctx, input, weight):
ctx.save_for_backward(input, weight)
input_temp, weight_temp = f(input, weight)
output = g(input_temp, weight_temp)
return output
def backward(ctx, grad_output):
input, weight = ctx.saved_tensors
'''
Compute backward with original inputs,
i.e., compute gradients of output directly to input and weight,
and ignore the functions in the middle (i.e., f(·), g(·))
'''
return grad_input, grad_weight
where f(·) and g(·) are custom functions for computing convolutions. One example would be to flatten the input, turning convolutions into matrix multiplications. In this case, f(·) could be unfold(·). As the backward path of unfold(·) is very slow, I don’t want this operation to be part of the computational graph. In other words, no matter what f(·) and g(·) are, I would like the autograd always treats this layer as output = nn.conv2d(input, weight), making the computational graph only include input, weight, and output without any middle variables and functions.
The pseudocode above requires customizing the gradient calculation of convolution. I tried several approaches, but none of them can reach the same speed as the original conv2d function. I’m wondering:
- How to implement the backward function of conv layer as fast as the original one in Pytorch?
- Otherwise, are there any other simple ways to handle this problem?
Please help and thanks!!