cuDNN-accelerated autoencoder: computing backward of cudnn_convolution.convolution_backward_input

username: Maxence_Ernoult
cooked: <p>Hi!</p>
<p>I have prototyped a convolutional autoencoder with two distinct sets of weights for the encoder (with parameters w_f) and for the decoder (w_b). I have naturally used <code>nn.Conv2d</code> and <code>nn.ConvTranspose2d</code> to build the encoder and decoder respectively. The rough context of study is on the one hand to learn w_f so that it minimizes a loss function defined at the output, and on the other hand to learn w_b so that it matches w_f^T. The forward pass typically reads as a convolution followed by a transpose convolution. So far, everything works out well.</p>
<p>Now I want to accelerate this code using the cuDNN-accelerated C++ functions <code>cudnn_convolution</code>, <code>cudnn_convolution_backward_weight</code> and  <code>cudnn_convolution_backward_input</code> (<a href="https://github.com/jordan-g/PyTorch-cuDNN-Convolution" class="inline-onebox" rel="noopener nofollow ugc">GitHub - jordan-g/PyTorch-cuDNN-Convolution: PyTorch extension enabling direct access to cuDNN-accelerated C++ convolution functions.</a>). I am writing my own conv layer (inheriting from <code>torch.autograd.Function</code>), defining its <code>forward</code> and <code>backward</code> methods by hand. In the forward method, we have a <code>cudnn_convolution</code> operation (parametrized by w_f) followed  by a <code>cudnn_convolution_backward_input</code> operation (parametrized by w_b). In the backward method, in order to compute the gradient of the loss with respect to w_b, I therefore need to backpropagate through <code>cudnn_convolution_backward_input</code>. I thought I could simply use <code>cudnn_convolution_backward_weight</code>, but it does not work. What is the good way to proceed to compute the gradient of w_b? I can provide the code of a concrete toy problem reproducing this issue if necessary. Many thanks for reading!</p>

username: albanD
cooked: <p>Hi,</p>
<p>That does sounds like the right strategy.<br>
In pseudo code, it would look like:</p>
<pre><code class="lang-auto">@staticmethod
def forward(ctx, inp, weight):
  ctx.save_for_backward(inp, weight)
  return cudnn_convolution(inp, weight)

@staticmethod
def backward(ctx, grad_out):
    inp, weight = ctx.saved_tensors
    return cudnn_convolution_backward_input(grad_out, weight), cudnn_convolution_backward_weight(grad_out, inp)
</code></pre>
<p>Doesn’t that work for you?</p>

username: Maxence_Ernoult
cooked: <p>Dear Alban,</p>
<p>Thank you very much for your time!<br>
Let me be more precise about the setting with the following pseudo-code:</p>
<pre><code class="lang-auto">    @staticmethod
    def forward(ctx, input, w_f, bias_f, w_b, stride, padding):
        
            y = cudnn_convolution.convolution(input, w_f, bias_f, 
                                              stride, padding,
                                              (1, 1), 1, False, False)
            
            r = cudnn_convolution.convolution_backward_input(input.shape, w_b.data, y, 
                                                             stride, padding, 
                                                             (1, 1), 1, False, False, False)
            
            ctx.save_for_backward(y, w_b)
            ctx.stride = stride
            ctx.padding = padding
            
            return r        
        
    @staticmethod
    def backward(ctx, grad_r):
        w_b_learning = ctx.w_b_learning
        stride, padding = ctx.stride, ctx.padding

        y, w_b = ctx.saved_tensors    
        
        grad_w_b = cudnn_convolution.convolution_backward_weight(..?..)
        grad_y = cudnn_convolution.convolution_backward_input(..?..)
        grad_w_f = cudnn_convolution.convolution_backward_weight(input, w_f.shape, grad_y, stride, padding, (1, 1), 1, False, False, False)
        grad_bias_f = torch.sum(grad_y, dim=[0, 2, 3]).squeeze(0)   
        
        return None, grad_w_f, grad_bias_f, grad_w_b, None, None
</code></pre>
<p>My question is therefore the following : how could I compute <code> grad_w_b</code> and <code>grad_y</code>? Can we also use <code>cudnn_convolution.convolution_backward_input</code> and <code>cudnn_convolution.convolution_backward_weight</code> with the “right” parameters? Or should we use another cpp function? It would be incredibly helpful, I have looked everywhere on the web and I could not find an answer :-(.</p>

username: albanD
cooked: <p>Why do you do the backward wrt to the input in the forward? That should only happen during the backward.</p>

username: Maxence_Ernoult
cooked: <p>Hi Alban,</p>
<p>I want to do in the forward pass something strictly equivalent to <code>ConvTranpose2d</code> (like we would do to define the forward pass of a convolutional autoencoder) but using these cuDNN accelerated functions. Which is why I thought in the first place to use <code>cudnn_convolution.convolution_backward_input</code> within the forward pass.</p>

username: Maxence_Ernoult
cooked: <p>To put it differently: based on the <code>cudnn_convolution.cpp</code> attached to his URL (<a href="https://github.com/ernoult/PyTorch-cuDNN-Convolution/blob/master/cudnn_convolution.cpp" class="inline-onebox" rel="noopener nofollow ugc">PyTorch-cuDNN-Convolution/cudnn_convolution.cpp at master · ernoult/PyTorch-cuDNN-Convolution · GitHub</a>), what should I add in this cpp file to define <code>cudnn_convolution_transpose</code>, <code>cudnn_convolution_transpose_backward_weight</code> and <code>cudnn_convolution_transpose_backward_input</code>? I should mention that I don’t know anything about cpp <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

username: albanD
cooked: <p>In that case, you don’t want to do the convolution itself.<br>
You can o that trick like we do it in pytorch core here: <a href="https://github.com/pytorch/pytorch/blob/a756a9e553ad23042fa39c6452d17aba72911abb/aten/src/ATen/native/cudnn/ConvShared.cpp#L368-L378" class="inline-onebox">pytorch/ConvShared.cpp at a756a9e553ad23042fa39c6452d17aba72911abb · pytorch/pytorch · GitHub</a></p>

username: Maxence_Ernoult
cooked: <p>Thank you so much Alban! This is exactly what I was looking for! I will try it out and let you know if it worked <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Maxence_Ernoult
cooked: <p>Hi again,</p>
<p>I am really sorry to bother you again, but I think I am failing to do something that must be relatively simple. Along the lines of the <code>cudnn_convolution.cpp</code> file of this repo ( <a href="https://github.com/ernoult/PyTorch-cuDNN-Convolution/blob/master/cudnn_convolution.cpp" rel="noopener nofollow ugc">PyTorch-cuDNN-Convolution/cudnn_convolution.cpp at master · ernoult/PyTorch-cuDNN-Convolution (github.com)</a>, I would like to create a wrapper around the C++ accelerated functions for <code>convolution_tranpose</code> and <code>convolution_transpose_backward_weight</code>. In <code>cudnn_convolution.cpp</code>, I have therefore added:</p>
<pre><code class="lang-auto">at::Tensor convolution_transpose(
    const at::Tensor&amp; input,
    const at::Tensor&amp; weight,
    const at::Tensor&amp; bias,
    c10::ArrayRef&lt;int64_t&gt; stride,
    c10::ArrayRef&lt;int64_t&gt; padding,
    c10::ArrayRef&lt;int64_t&gt; output_padding,
    c10::ArrayRef&lt;int64_t&gt; dilation,
    int64_t groups,
    bool benchmark,
    bool deterministic) {

    return at::cudnn_convolution_transpose(
        input,
        weight,
        bias,
        padding,
        output_padding,
        stride,
        dilation,
        groups,
        benchmark,
        deterministic); 
}


at::Tensor convolution_transpose_backward_weight(
    const at::Tensor&amp; input,
    c10::ArrayRef&lt;int64_t&gt; weight_size,
    const at::Tensor&amp; grad_output,
    c10::ArrayRef&lt;int64_t&gt; stride,
    c10::ArrayRef&lt;int64_t&gt; padding,
    c10::ArrayRef&lt;int64_t&gt; dilation,
    int64_t groups,
    bool benchmark,
    bool deterministic,
    bool allow_tf32) {

    return at::cudnn_convolution_transpose_backward_weight(
        weight_size,
        grad_output,
        input,
        padding,
        stride,
        dilation,
        groups,
        benchmark,
        deterministic,
        allow_tf32);
}
</code></pre>
<p>Then, I load this cpp modules from a python script through:</p>
<pre><code class="lang-auto">cudnn_convolution = load(name="cudnn_convolution", sources=["cudnn_convolution.cpp"], verbose=True)
</code></pre>
<p>Finally, I try to use transpose convolution with the following script:</p>
<pre><code class="lang-auto">weight = torch.zeros(64, 3, 5, 5).to('cuda')
bias   = torch.zeros(64).to('cuda')

stride   = (2, 2)
padding  = (0, 0)
output_padding  = (0, 0)
dilation = (1, 1)
groups   = 1
dummy_input = torch.zeros(128, 64, 14, 14).to('cuda')
output = cudnn_convolution.convolution_transpose(dummy_input, weight, bias, stride, padding, output_padding, dilation, groups, False, False)
</code></pre>
<p>but it throws the following error message:</p>
<pre><code class="lang-auto">AttributeError: module 'cudnn_convolution' has no attribute 'convolution_transpose'
</code></pre>
<p>Do you have any idea of what happens? I am a bit desperate <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>