Manually/explicitly calculate gradients of Conv kernels

111414: <p><strong>1. Background:</strong></p>
<p>I can calculate the gradient of <code>x</code> with respect to a cost function <code>loss</code> in two ways: (1) manually writing out the explicit and analytic formula, and (2) using <code>torch.autograd</code> package. Here is my example:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as F

for i in range(10):
    x = torch.randn(8, 1, 128, 128)
    y = torch.randn(8, 512, 4, 4)
    k = torch.randn(512, 1, 32, 32)

    loss = lambda z: 0.5 * (F.conv2d(z, k, stride=32) - y).pow(2).sum(dim=[1,2,3])  # cost function is [(1/2)||k*x-y||_F^2]

    # 1: calculate gradient of x explicitly and manually
    x_grad_manual = F.conv2d(x, k, stride=32) - y
    x_grad_manual = F.conv_transpose2d(x_grad_manual, k, stride=32)

    # 2: calculate gradient of x using torch.autograd
    x_var = torch.autograd.Variable(x, requires_grad=True)
    x_var_loss = loss(x_var)
    x_grad_auto = torch.autograd.grad(x_var_loss, x_var, torch.ones_like(x_var_loss), create_graph=True, retain_graph=True)[0]

    # check if the results of implementations 1 and 2 are equal
    print((x_grad_manual - x_grad_auto).pow(2).mean())
</code></pre>
<p>Since the mean squared errors of the results of the above two implementations are very small (about 3.4*10^(-8)), I think that their should be mutually matched and the manual implementation works correctly.</p>
<p><strong>2. My Problem</strong></p>
<p>I am confused by how to explicitly write out the gradients of variables (features and Conv kernels) conveniently with some compound processes? For instance, I do not know how to calculate the gradients of feature <code>x</code> and Conv kernel <code>w1</code> in the following context:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as F

for i in range(10):
    x = torch.randn(8, 32, 128, 128)
    y = torch.randn(8, 512, 4, 4)
    k = torch.randn(512, 1, 32, 32)
    w1 = torch.randn(1, 32, 3, 3)

    def loss(z, w):
        z_forward = F.conv2d(z, w, padding=1)  # z = w1 * x
        return 0.5 * (F.conv2d(z_forward, k, stride=32) - y).pow(2).sum(dim=[1,2,3])  # cost function is [(1/2)||k*z-y||_F^2]

    # 1: calculate gradients of x and w1 explicitly and manually
    x_grad_manual = ???
    w1_grad_manual = ???

    # 2: calculate gradients of x and w1 using torch.autograd
    x_var = torch.autograd.Variable(x, requires_grad=True)
    w1_var = torch.autograd.Variable(w1, requires_grad=True)
    x_var_loss = loss(x_var, w1_var)
    x_grad_auto, w1_grad_auto = torch.autograd.grad(x_var_loss, [x_var, w1_var], torch.ones_like(x_var_loss), create_graph=True, retain_graph=True)

    # check if the results of implementations 1 and 2 are equal
    print((x_grad_manual - x_grad_auto).pow(2).mean())
    print((w1_grad_manual - w1_grad_auto).pow(2).mean())
</code></pre>
<p><strong>3. Extension:</strong></p>
<p>Furthermore, if the forwarding process is more complicated than the above one, with two middle Conv layers and a <code>ReLU</code> activation, how can I write out the gradients? Please see the following problem:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as F

for i in range(10):
    x = torch.randn(8, 32, 128, 128)
    y = torch.randn(8, 512, 4, 4)
    k = torch.randn(512, 1, 32, 32)
    w1 = torch.randn(32, 32, 3, 3)
    w2 = torch.randn(1, 32, 3, 3)

    def loss(z, q1, q2):
        z_forward = F.conv2d(z, q1, padding=1)  # z = w1 * x
        z_forward = F.relu(z_forward, inplace=True)  # z = ReLU(w1 * x)
        z_forward = F.conv2d(z_forward, q2, padding=1)  # z = w2 * ReLU(w1 * x)
        return 0.5 * (F.conv2d(z_forward, k, stride=32) - y).pow(2).sum(dim=[1,2,3])  # cost function is [(1/2)||k*z-y||_F^2]

    # 1: calculate gradients of x, w1 and w2 explicitly and manually
    x_grad_manual = ???
    w1_grad_manual = ???
    w2_grad_manual = ???

    # 2: calculate gradients of x, w1 and w2 using torch.autograd
    x_var = torch.autograd.Variable(x, requires_grad=True)
    w1_var = torch.autograd.Variable(w1, requires_grad=True)
    w2_var = torch.autograd.Variable(w2, requires_grad=True)
    x_var_loss = loss(x_var, w1_var, w2_var)
    x_grad_auto, w1_grad_auto, w2_grad_auto = torch.autograd.grad(x_var_loss, [x_var, w1_var, w1_var], torch.ones_like(x_var_loss), create_graph=True, retain_graph=True)

    # check if the results of implementations 1 and 2 are equal
    print((x_grad_manual - x_grad_auto).pow(2).mean())
    print((w1_grad_manual - w1_grad_auto).pow(2).mean())
    print((w2_grad_manual - w2_grad_auto).pow(2).mean())
</code></pre>
<p><strong>4. Guarantee of Differentiability</strong></p>
<p>Like my first example, I hope that the manual gradient calculations are totally explicit and themselves are differential, such that I can inject some of the processes in my neural network implementation. Could you please teach me how to achieve this?</p>
<p><strong>5. The Reason of Posting This Problem</strong></p>
<p>In a neural network I constructed, it is needed to calculate the gradients of some features and Conv kernels with respect to my pre-defined cost functions (as you can see above). In my current implementations, I directly employ <code>torch.autograd</code> package to calculate various gradients. However, it seems that there are some mistakes accumulated which misleads the learning process when I train such a neural network.</p>
<p>(The whole neural network has its own <code>loss</code> function and <code>backward</code> process. I just added some extra inner gradient calculations to achieve my goals.)</p>
<p>I conjecture that I should calculate the gradients manually and not directly use <code>torch.autograd</code> in a common network forwarding process, since some computational graphs and backwards may be nested and lead to the wrong weight updates.</p>
<p>In my experiments, I train two networks (with manual and auto-calculations, like the first example) and get similar results. But when I extend to more complicated forwardings (like my posted two problems), the training processes would not be stable. So I want to manually write out the gradients to avoid the implementation mistakes and conduct more experiements.</p>

KFrank: <p>Hi Bin!</p>
<aside class="quote no-group" data-username="111414" data-post="1" data-topic="150462" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/111414/48/31166_2.png" class="avatar"> 111414:</div>
<blockquote>
<p>I am confused by how to explicitly write out the gradients of variables (features and Conv kernels) conveniently with some compound processes?</p>
</blockquote>
</aside>
<p>To get the gradient of the result of one function applied to the result of<br>
another function, that is, of the <em>composition</em> of two functions, you would<br>
use the <a href="https://en.wikipedia.org/wiki/Chain_rule" rel="noopener nofollow ugc">chain rule.</a>  This is how autograd computes the gradient when<br>
many functions are composed together, such as the successive layers<br>
in a network.</p>
<aside class="quote no-group">
<blockquote>
<p>I directly employ <code>torch.autograd</code> package to calculate various gradients. However, it seems that there are some mistakes accumulated which misleads the learning process when I train such a neural network.<br>
…<br>
I conjecture that I should calculate the gradients manually and not directly use <code>torch.autograd</code> in a common network forwarding process,</p>
</blockquote>
</aside>
<p>It is true that floating-point round-off error can accumulate during<br>
backpropagation (as it can during the forward pass, as well).  Underflow<br>
and overflow "errors’ can occur as well.  Nonetheless, autograd does<br>
an altogether solid job of performing these numerical computations.<br>
It is unlikely that you will be able to do better calculating your own<br>
gradients, or, in effect, writing your own version of autograd.</p>
<p>If you really are having problems with numerical stability during<br>
backpropagation, you would be better off identifying the root cause<br>
and addressing it directly, presumably by using more numerically<br>
stable functions or implementations in your forward pass.</p>
<p>Best.</p>
<p>K. Frank</p>