A strange backward results of my max pooling CUDA extension

username: _ju
cooked: <p>I want to implement a max pooling function and I come across a strange bug. I simplify my codes and reproduce the problem as follow:</p>
<p>MaxPooling.py:</p>
<pre><code class="lang-auto"># MaxPooling.py

class MaxPoolFunction(Function):

    @staticmethod
    def forward(ctx, features):
        n_features, n_channels = features.size()
        output = torch.zeros(n_channels, device=features.device, dtype=torch.float)
        argmax = torch.empty(n_channels, device=features.device, dtype=torch.long).fill_(-1)
        max_pooling.forward(features, output, argmax)
        ctx.argmax = argmax
        ctx.n_features = n_features
        ctx.n_channels = n_channels
        return output

    @staticmethod
    def backward(ctx, grad_output):
        print('grad_output:', grad_output, '\n')

        grad_input = torch.zeros(ctx.n_features, ctx.n_channels, device=grad_output.device, dtype=grad_output.dtype)
        # if I use original grad_output as argument, the result will be wrong.
        max_pooling.backward(grad_output, ctx.argmax, grad_input)

        # But if I make a copy of grad_output as argument, the result will be right.
        new_grad_output = grad_output.clone().detach()
        assert torch.sum(grad_output != new_grad_output) == 0
        grad_input_1 = torch.zeros(ctx.n_features, ctx.n_channels, device=grad_output.device, dtype=grad_output.dtype)
        max_pooling.backward(new_grad_output, ctx.argmax, grad_input_1)

        print("grad_input: ", grad_input)
        print("grad_input_1: ", grad_input_1)

        return grad_input


# test
if __name__ == '__main__':
    features = torch.rand((2, 8)).cuda().requires_grad_()
    result = MaxPoolFunction.apply(features)
    print("features:", features)
    print("max pooling results:", result)
    torch.sum(result).backward()
</code></pre>
<p>max_pooling.cpp:</p>
<pre><code class="lang-auto">// max_pooling.cpp

void forwardLaucher(int n_channels, int n_features, float* features, float* output, long* argmax);
void forward(at::Tensor features, at::Tensor output, at::Tensor argmax)
{
    int n_features = features.size(0);
    int n_channels = features.size(1);
    forwardLaucher(n_channels, n_features, features.data&lt;float&gt;(), output.data&lt;float&gt;(), argmax.data&lt;long&gt;());
}

void backwardLaucher(int n_channels, float* grad, long* argmax, float* output);
void backward(at::Tensor grad, at::Tensor argmax, at::Tensor output)
{
    int n_channels = grad.size(0);
    backwardLaucher(n_channels,  grad.data&lt;float&gt;(), argmax.data&lt;long&gt;(), output.data&lt;float&gt;());
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("forward", &amp;forward, "forward");
    m.def("backward", &amp;backward, "backward");
}
</code></pre>
<p>max_pooling_kernel.cu:</p>
<pre><code class="lang-auto">// max_pooling_kernel.cu

__global__ void forwardKernel(int n_channels, int n_features, float* features, float* output, long* argmax)
{
    int c = threadIdx.x;
    for(int i=0; i&lt;n_features; i++)
    {
        if(argmax[c]==-1 || output[c]&lt;features[i*n_channels+c])
        {
            output[c] = features[i*n_channels+c];
            argmax[c] = i;
        }
    }
}
void forwardLaucher(int n_channels, int n_features, float* features, float* output, long* argmax)
{
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    forwardKernel&lt;&lt;&lt;1, n_channels, 0, stream&gt;&gt;&gt;(n_channels, n_features, features, output, argmax);
}


__global__ void backwardKernel(int n_channels, float* grad, long* argmax, float* output)
{
    int c = threadIdx.x;
    long idx = argmax[c];
    if(idx != -1)
    {
        atomicAdd(&amp;output[idx*n_channels+c], grad[c]);
    }
}

void backwardLaucher(int n_channels, float* grad, long* argmax, float* output)
{
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    backwardKernel&lt;&lt;&lt;1, n_channels, 0, stream&gt;&gt;&gt;(n_channels, grad, argmax, output);
}
</code></pre>
<p>When I run MaxPooling.py, I get following results:</p>
<blockquote>
<p>features: tensor([[0.1430, 0.9068, 0.1484, 0.9407, 0.7852, 0.3469, 0.0498, 0.9610],<br>
[0.1713, 0.9893, 0.2032, 0.9011, 0.1194, 0.1983, 0.1744, 0.2393]],<br>
device=‘cuda:0’, requires_grad=True)</p>
<p>max pooling results: tensor([0.1713, 0.9893, 0.2032, 0.9407, 0.7852, 0.3469, 0.1744, 0.9610],<br>
device=‘cuda:0’, grad_fn=)</p>
<p>grad_output:  tensor([1., 1., 1., 1., 1., 1., 1., 1.], device=‘cuda:0’)</p>
<p>grad_input:  tensor([[0.0000, 0.0000, 0.0000, 0.9407, 0.7852, 0.3469, 0.0000, 0.9610],<br>
[1.0000, 1.5463, 0.2032, 0.0000, 0.0000, 0.0000, 0.1744, 0.0000]],<br>
device=‘cuda:0’)</p>
<p>grad_input_1:  tensor([[0., 0., 0., 1., 1., 1., 0., 1.],<br>
[1., 1., 1., 0., 0., 0., 1., 0.]], device=‘cuda:0’)</p>
</blockquote>
<p>We can see that the <code>grad_output</code> in backward function is same as <code>new_grad_output</code>. But using the original <code>grad_output</code> get wrong backward result. Why?</p>
<p>The codes should have limited that the grad_input only depends on the grad_output and argmax. But from the results we can also see that some values of grad_input are the values of max pooling results. It’s very strange.</p>

username: albanD
cooked: <p>Hi,</p>
<p>In you <code>backward</code> cpp method, you do <code>int n_channels = grad.size(0);</code> while to be consistent with the rest of your code it should be <code>int n_channels = grad.size(1);</code> no? Or do you expect some transpose somewhere?</p>

username: _ju
cooked: <p>Thanks for your reply.</p>
<p>I think <code>int n_channels = grad.size(0)</code> is right. Because the of input of max pooling has size (2, 8). Here 2 denotes n_features and 8 denotes n_channels. And my  forward CUDA method do max pooling along dimension 0. So the output of max pooling has size (8,) , and the grad in backward method also has size (8,). So I do <code>int n_channels = grad.size(0)</code>.</p>
<p>I think my question is not so clear.<br>
In normal way to implement a backward methods of max pooling, I should have created a new tensor with size (2, 8) , and element-wise add the <code>grad</code> to this new tensor according to the <code>argmax</code>. So my codes do.</p>
<pre><code class="lang-auto">grad_input = torch.zeros(ctx.n_features, ctx.n_channels, device=grad_output.device, dtype=grad_output.dtype)
max_pooling.backward(grad_output, ctx.argmax, grad_input)
print('grad_input': grad_input)
</code></pre>
<p>But I get an unexpected results when I printed the grad_input after max_pooling.backward():</p>
<blockquote>
<p>grad_input: tensor([[0.0000, 0.0000, 0.0000, 0.9407, 0.7852, 0.3469, 0.0000, 0.9610],<br>
[1.0000, 1.5463, 0.2032, 0.0000, 0.0000, 0.0000, 0.1744, 0.0000]],<br>
device=‘cuda:0’)</p>
</blockquote>
<p>the argument <code>grad_output</code> in <code>max_pooling.backward(grad_output, ctx.argmax, grad_input)</code> is a all-ones tensor. (I printed it). So why the backward methods produce float number <code>0.9407 0.7852 0.3469 ...</code> ?</p>
<p>Anthor strange phenomenon is that, if I make a copy of original <code>grad_out</code> as argument for backward method. The result will be right!</p>
<pre><code class="lang-auto">new_grad_output = grad_output.clone().detach()
grad_input_1 = torch.zeros(ctx.n_features, ctx.n_channels, device=grad_output.device, dtype=grad_output.dtype)
max_pooling.backward(new_grad_output, ctx.argmax, grad_input_1)
print('grad_input_1': grad_input_1)
</code></pre>
<blockquote>
<p>grad_input_1: tensor([[0., 0., 0., 1., 1., 1., 0., 1.],<br>
[1., 1., 1., 0., 0., 0., 1., 0.]], device=‘cuda:0’)</p>
</blockquote>
<p>Are there any differences bettween original <code>grad_out</code> and copied  <code>grad_out</code>, leading to different results?</p>

username: _ju
cooked: <p>I have solved the problem.</p>
<p>I used <code>grad_output.contiguous()</code> before <code>max_pooling.backward(grad_output, ctx.argmax, grad_input)</code> in py file, and it worked.</p>
<pre><code class="lang-auto">    @staticmethod
    def backward(ctx, grad_output):
        # grad_output is not contiguous here. But I do not know why.
        grad_output = grad_output.contiguous()
        grad_input = torch.zeros(ctx.n_features, ctx.n_channels, device=grad_output.device, dtype=grad_output.dtype)
        max_pooling.backward(grad_output, ctx.argmax, grad_input)
        return grad_input
</code></pre>
<p>Thanks again.</p>

username: ryanbty
cooked: <p>Hello, nice topic !</p>
<p>Do you know if it is possible to change the value of the derivative of maxpooling in the backward fonction ? (<a class="mention" href="/u/_ju">@_ju</a> and <a class="mention" href="/u/alband">@albanD</a>)</p>
<p>For example, if F is maxpool1d and x = (1,1,1) then F’(x) = (1/3,1/3,1/3).</p>
<p>class MyMaxPool1d(torch.autograd.Function):</p>
<pre><code>@staticmethod
def forward(ctx, input, kernel_size): 
    """
    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.
""" 
    
    ctx.save_for_backward(input)
    m = nn.MaxPool1d(kernel_size, stride=1)
    output = m(input)
    return 3*output.sum()

@staticmethod
def backward(ctx, grad_output):
    
    input, = ctx.saved_tensors
    grad_input = grad_output.clone()
    return 2*grad_input, None
</code></pre>
<p>if <strong>name</strong> == ‘<strong>main</strong>’:<br>
features =  torch.ones(1, 4, 4,requires_grad=True)<br>
result = MyMaxPool1d.apply(features,3)<br>
print(“features:”, features)<br>
result.backward()</p>
<p>I have that error : RuntimeError: Function MyMaxPool1dBackward returned an invalid gradient at index 0 - got [] but expected shape compatible with [1, 4, 4]</p>

username: albanD
cooked: <p>Hi,</p>
<p>This is indeed the right way to do this.<br>
You need to ensure that the grad_input is the same size as the input to the forward function though</p>