Xla: Gradient Penalty with stride=2 Convolution would crash

[luzai]

🐛 Bug

To Reproduce

Gradient penalty requires high-order gradient of a loss, w.r.t the weights, and is important in research areas, like meta learning. May I ask why the code below would crash the colab session?

Note that the penalty in the code below is different from the gradient penalty/input regularization in adversarial training or wgan, it is w.r.t the weight in convolution layer. However, this features is required in meta learning tasks.

import torch
import torch_xla
import torch_xla.core.xla_model as xm
from torch import nn 
import torch.nn.functional as F

stride = 2 # set stride=2 will cause the colab session crashed for an unknown reason
device = xm.xla_device()
data = torch.rand(32,3,32,32).to(device)
conv = nn.Conv2d(3,20,3, stride=stride, padding=1, bias=False).to(device) 
loss = F.adaptive_avg_pool2d(conv(data), 1).view(32,20).mean()
dw = torch.autograd.grad(loss, conv.parameters(), create_graph=True, retain_graph=True, only_inputs=True ) [0]
(dw*dw).mean().backward() 

Steps to reproduce the behavior:

1. choose the nightly version of torch_xla. And backend of TPU
2. run the code snippet in google colab session
3. The session would pop a message: "Your session crashed for an unknown reason."

Expected behavior

When stride=2, The session would pop a message: "Your session crashed for an unknown reason."
When stride=1, the code runs successfully.
The code can run successfully on gpu, with device='cuda:0'.

Environment

• Reproducible on XLA backend [CPU/TPU]: TPU
• torch_xla version: nightly

[dlibenzi]

The real error is:

Invalid argument: conv_backward_input: Size of out_backprop doesn't match computed: actual = 3, computed = 4 spatial_dim: 2 input: 32 filter: 16 output: 3 stride: 1 dilation: 2 vs. OK

https://gist.github.com/dlibenzi/c5333bbf854cfcbed872ed7b1eed1cf2

[luzai]

Thank you for your response!
I think backpropogation of conv with stride=2 involves conv with dilation=2 on the input gradient. I am wondering why the size tensor would mismatch. Does the error means the expected out_backprop should be size=3, but computed size is 4?

[dlibenzi]

There is some bug in the computations. I have not had time to look into it.

[JackCaoG]

The problems seems to come from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp#L1120. For the context, pt/xla overwrites two pytorch native op convolution_overrideable and convolution_backward_overrideable. Pytorch _convolution_double_backward calls at::_convolution during the backward with Input shape [3,20,3,3], mask shape [20, 32, 16, 16], dilation = [2,2] and transposed=true.

In our implemantion, transposed convolution forward calls BuildConvBackwardInput and then we have

Input shape[3, 32, 32, 32] -> [32,32] is calculated based on grad_output shape and kernel shape, [3,32] is taken from
grad_output shape and kernel shape directly. 3 = batch size, 32 = input feature

kernel shape [16, 16, 32, 20, ] -> transposed from [20, 32, 16, 16]
grad_output shape [3, 20, 3, 3, ]
stride[1, 1]
dilation[2, 2]
padding[1, 1]

The failure was that for a single depth image of [32, 32] and filter of size [16, 16] with dilation [2,2] and padding [1,1], the output should be [4,4]. However the grad_output is of size [3, 3]. I need to figure out where is this stride->dilation comes from.

[JackCaoG]

Ah, ok. The swapping is from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp#L1084. Testing with pytorch native code to check how is this handled.

[JackCaoG]

I chatted with @ailzhang and we figure out that there is a bug in our transposed_convolution_forward. In short, we need to pad the input tensor at spatial dimension with the output_padding at the end before performing the convolution. transposed_convolution_forward is essentially the same as transposed_convolution_backward.

In the example you provided,

Conv forward
Input,             weight,           res
[32,3,32,32]    [20, 3, 3,3]     [32,20,16,16]   stride = [2,2]  dilation = [1, 1] padding[1,1] 

Conv backward
Grad_output     kernel           input_shape           (treated as backward)
Input           kernel            res                  (treated as transpose forward)
[3,20,3,3]    [20,32,16,16]     [3, 32, 32, 32] (computed)       stride [1,1 ]  dilation[2,2], padding[1,1], output padding [1,1]

output padding is something calculated by pytorch and passed to the pt/xla. We currently used it to computed the output side but we also need to pad our input. Note that if the kernel is of size [20, 3, 4, 4], the result of the convolution is still of size 32,20,16,16 but during the backward the output_padding will be [0, 0]. If you are interested, you can look at the output padding definition here.

Another good explaination would be

The padding argument effectively adds dilation * (kernel_size - 1) - padding amount of zero padding to both sizes of the input. This is set so that when a Conv2d and a ConvTranspose2d are initialized with same parameters, they are inverses of each other in regard to the input and output shapes. However, when stride > 1, Conv2d maps multiple input shapes to the same output shape. output_padding is provided to resolve this ambiguity by effectively increasing the calculated output shape on one side. Note that output_padding is only used to find output shape, but does not actually add zero-padding to output.

I will work on a formal fix and test tomorrow

[luzai]

Thank you very much for the detailed and professional explanation!
I get to understand that output_padding is used to resolve the ambiguity brought by stride>1. Thank you for maintaining this complex system, and looking forward to the fix.

[JackCaoG]

This turned out to be a bit more complicated than I thought. In the transposed_conv, conv_backward, pass, we can simply pad the input and the result size won't be affected. However in the transposed_conv_backward_bias or transposed_conv_backward_kernel phase test is reporting an incompatible shape error. I think we need to unpad the grad_output somewhere, need to think about it a bit more about this.. The good news is that after some modification our unit test is doing a great job catching all of these corner cases!

[JackCaoG]

Close this bug as the fix pr has merged