Access cuDNN convolution grad input and grad weight computation function

flauron: <p>Hi,</p>
<p>I am currently working on a problem of gradient pruning.</p>
<p>In the context of convolution computation, during the backpropagation, I need to alter the gradient of the weight (grad_weight) computation  by slicing the gradient of the output (grad_output) before doing the actual  gradient computation.</p>
<p>So to achieve this goal I need acces to the functions performing backpropagarion computation for convolution.</p>
<p>I don’t want to use the functions <strong>conv2d_input</strong> and  <strong>conv2d_weight</strong> because they are slow. Binding the function <strong>convolution_backward</strong> does not solve my problem since I also compute the gradient of the input (grad_input) without slincing grad_output.</p>
<p>I want to use  CUDA/cuDNN function to do the job.</p>
<p>My best guess so far was this example on github:</p><aside class="onebox githubfolder" data-onebox-src="https://github.com/jordan-g/PyTorch-cuDNN-Convolution/tree/master">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/jordan-g/PyTorch-cuDNN-Convolution/tree/master" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h3><a href="https://github.com/jordan-g/PyTorch-cuDNN-Convolution/tree/master" target="_blank" rel="noopener nofollow ugc">GitHub - jordan-g/PyTorch-cuDNN-Convolution: PyTorch extension enabling direct...</a></h3>

  <p><a href="https://github.com/jordan-g/PyTorch-cuDNN-Convolution/tree/master" target="_blank" rel="noopener nofollow ugc">master</a></p>

  <p><span class="label1">PyTorch extension enabling direct access to cuDNN-accelerated C++ convolution functions. - jordan-g/PyTorch-cuDNN-Convolution</span></p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>This example uses <strong>cudnn_convolution_backward_weight</strong> and <strong>cudnn_convolution_backward_input</strong> as function to compute the associated gradient.</p>
<p>But, from what I understand theses functions have been removed from the ATen API (for code rule compliance)and are no longer accessible (they are not anymore written inside native_functions.yaml)</p>
<p>My point is I really really need to have a fine control over the grad_input and grad_weight (and grad_bias) computation, So how I can do that?</p>
<p>More generaly, how to make accessible for python/torch binding functions that are not listed inside native_functions.yaml?</p>
<p>Sorry for the long post.<br>
Thanks in advance for your answers and guidance.</p>

ptrblck: <aside class="quote no-group" data-username="flauron" data-post="1" data-topic="214064">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/f/ba8739/48.png" class="avatar"> flauron:</div>
<blockquote>
<p>I don’t want to use the functions <strong>conv2d_input</strong> and <strong>conv2d_weight</strong> because they are slow.</p>
</blockquote>
</aside>
<p>These methods will still dispatch to <code>cudnn</code> as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">device = "cuda"
x = torch.randn(1, 3, 24, 24, requires_grad=True, device=device)
w = torch.randn(3, 3, 3, 3, requires_grad=True, device=device)
grad_output = torch.randn_like(x)


activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
sort_by_keyword = device + "_time_total"

with profile(activities=activities, record_shapes=True) as prof:
    wgrad = torch.nn.grad.conv2d_weight(x, w.shape, grad_output, stride=1, padding=1)
print(prof.key_averages().table(sort_by=sort_by_keyword, row_limit=10))
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                                    Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                              aten::convolution_backward        12.68%     212.324us        15.07%     252.310us     252.310us       4.608us       100.00%       4.608us       4.608us             1  
# void cudnn::cnn::wgrad2d_grouped_direct_kernel&lt;false...         0.00%       0.000us         0.00%       0.000us       0.000us       4.608us       100.00%       4.608us       4.608us             1  
#                                         aten::new_empty         8.46%     141.694us        79.58%       1.332ms       1.332ms       0.000us         0.00%       0.000us       0.000us             1  
#                                             aten::empty        71.51%       1.197ms        71.51%       1.197ms     598.527us       0.000us         0.00%       0.000us       0.000us             2  
#                                            aten::expand         0.79%      13.165us         1.06%      17.703us      17.703us       0.000us         0.00%       0.000us       0.000us             1  
#                                        aten::as_strided         0.27%       4.538us         0.27%       4.538us       4.538us       0.000us         0.00%       0.000us       0.000us             1  
#                                                aten::to         0.10%       1.703us         0.10%       1.703us       1.703us       0.000us         0.00%       0.000us       0.000us             1  
#                                         cudaEventRecord         0.34%       5.731us         0.34%       5.731us       5.731us       0.000us         0.00%       0.000us       0.000us             1  
#                                   cudaStreamIsCapturing         0.06%       1.022us         0.06%       1.022us       1.022us       0.000us         0.00%       0.000us       0.000us             1  
#                                   cudaStreamGetPriority         0.06%       1.072us         0.06%       1.072us       1.072us       0.000us         0.00%       0.000us       0.000us             1  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
# Self CPU time total: 1.674ms
# Self CUDA time total: 4.608us
</code></pre>
<p>so the change in the interface should not affect the runtime.</p>

flauron: <p>Thanks a lot.</p>
<p>Your answer help me to understand a few things. Correct me if I 'm wrong.</p>
<p>For exemple <strong>conv2d_input</strong> call the function <strong>convolution_backward</strong> through this call:<br>
(In pytorch/torch/nn/grap.py)</p>
<pre data-code-wrap="python"><code class="lang-python">torch.ops.aten.convolution_backward(
        grad_output,
        input,
        weight,
        None,
        _pair(stride),
        _pair(padding),
        _pair(dilation),
        False,
        [0],
        groups,
        (True, False, False),
    )[0]
</code></pre>
<p>which is the main function for computing backward convolution, and I can call it this way , because it has been listed in the file <em>pytorch/aten/src/ATen/native/native_functions.yaml</em>.</p>
<p>So basically I can also call a function, listed in this file, for computing forward pass for batch normalisation:<br>
torch.ops.aten._batch_norm_impl_index<br>
Which gives me an <code>output</code> tensor , <code>save_mean</code> tensor, <code>save_var</code> tensor, <code>reserve</code> tensor (and a fifth one not usefull)<br>
And compute the backward pass in a similar way it is done for the convolution with this function, using <code>output</code> , <code>save_mean</code>,<code>save_var</code>, and <code>reserve</code> .</p>
<p>(e.g for grad_input)</p>
<pre data-code-wrap="python"><code class="lang-python">grad_input = torch.ops.aten.batch_norm_backward(
                grad_output,
                input,
                weight,
                running_mean,
                running_var,
                save_mean,
                save_var, False, eps,
                (True, False, False), reserve)[0]
</code></pre>
<p>Anyways thrank you.</p>