So far, I’ve check all these posts that I found on the issue, but None
of them seems to fit my case.
- backward-not-called-bugs
- backward-not-called
- custom-autograd-function-backward-pass-not-called
- using-backward-of-custom-autograd-function-with-loss-backward
Here is my code
class _Conv2dCustom(autograd.Function):
# Kernel has to be of odd size
@staticmethod
def forward(ctx, input, weight, bias=None, padding=0):
print('########## _Conv2dCustom.forward')
input = input.clone()
ctx.save_for_backward(input, weight, bias)
padding = (padding, padding) if type(padding) is int else padding
# input = input.clone()
output = _Conv2dCustom._conv2d(input, weight, padding)
if bias is not None:
output += bias.view(-1, 1, 1)
return output
@staticmethod
def backward(ctx, grad_output):
print('########## _Conv2dCustom.backward')
# Kernel has to be of odd size
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
if ctx.needs_input_grad[0]:
# With padding, conv produce same size grad_output
padding = weight.size(-1)//2
# Difference to produce same size as input
padding += (grad_output.size(-1)-input.size(-1))//2
weight_t = torch.transpose(weight, -1, -2)
grad_input = _Conv2dCustom._conv2d(grad_output, weight_t, padding)
if ctx.needs_input_grad[1]:
grad_weight = _Conv2dCustom._conv2d(input, grad_output, padding_used)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum([0, 2, 3])
return grad_input, grad_weight, grad_bias
@staticmethod
def _conv2d(input, weight, padding=(0, 0), stride=(1, 1)):
output_h = (input.size(-2) - weight.size(-2) + 2*padding[0])//stride[0] + 1
output_w = (input.size(-1) - weight.size(-1) + 2*padding[1])//stride[1] + 1
Xunfold = F.unfold(input, weight.size()[2:], padding=padding, stride=stride)
weight_flat = weight.view(weight.size(0), -1)
conv = weight_flat@Xunfold
output = conv.view(input.size(0), weight.size(0), output_h, output_w)
return output
class Conv2dCustom(nn.Conv2d):
# Torch master as in Feb 23rd 2020.
def _conv_forward(self, *args, **kwargs):
raise NotImplementedError('This Method Is Not Implemented')
# Torch v1.4.0
def conv2d_forward(self, input, weight):
conv2d = _Conv2dCustom.apply
conv = conv2d(input, weight, self.bias, self.padding)
return conv
Output
==========
Started Training
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
########## _Conv2dCustom.forward
==========
Finished Training
In my Network, I have two of these Conv2dCustom
layers, and I also put hooks on each layer. I noticed that the backward hook for all my layers, except the most inner Conv2dCustom
layer.