Backward of non-sequential module

username: vuhoangminh
cooked: <p>Hi everyone,</p>
<p>I would like to implement a module that I can perform some operations to intermediate layers such that:</p>
<p>(1) <code>requires_gradient = False</code> for these operations <code>a, b</code></p>
<p>(2) <code>requires_gradient = True</code> for <code>out</code></p>
<p>(3) force PyTorch to build Autograd graph such that the order is: <code>input -&gt; out (conv1) -&gt; out (relu1)</code>. It means that PyTorch won’t need to take care of computing gradients for any <code>function_a</code> and <code>function_b</code> in the backward pass. The reason is: <code>function_a</code> and <code>function_b</code> were not provided by PyTorch, and it is very hard to compute their derivatives.</p>
<p>I am not sure whether it is feasible or not. Do you have any ideas how to proceed?</p>
<p>Thank you so much for reading.</p>
<p>A simple network can be seen as follows:</p>
<pre><code class="lang-auto">class CustomModule(nn.Module):
    def __init__(self, module):
        super(CustomModule, self).__init__()
        self.module = module

    def forward(self, input):
        a = function_a(input)       # no need gradient in backward
        out = self.module(input)    # need gradient in backward
        b = function_b(out)         # no need gradient in backward
        return b


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = CustomModule(nn.Conv3d(1, 100, kernel_size=1))
        self.relu1 = CustomModule(nn.ReLU())

    def forward(self, input):
        conv1 = self.conv1(input)
        relu1 = self.relu1(conv1)
        return relu1
</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p>The gradients are computed using backpropagation. So to get the gradients for a module, you need all use of its output to be done by differentiable functions.<br>
If you do:</p>
<pre><code class="lang-auto">a = func1(input)
b = module(a)
c = func2(b)
loss = crit(c)
</code></pre>
<p>And your learnable parameters are in module.<br>
You will need <code>func2</code> and <code>crit</code> to be differentiable. But not <code>func1</code>.</p>

username: vuhoangminh
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,</p>
<p>Thank you so much for your response.</p>
<p>Is it true that if my learnable parameters are in <code>module1</code> and <code>module2</code>, and I do:</p>
<pre><code class="lang-auto">a1 = func1(input)
b1 = module1(a1)
c1 = func2(b1)

a2 = func1(c1)
b2 = module2(a2)
c2 = func2(b2)

loss = crit(c2)
</code></pre>
<p>(1) I will need only <code>crit</code>and <code>func2</code> to be differentiable? How does backward-pass work in the aforementioned scenario?</p>
<p>(2) As I understand, my problem is a kind of disconnected graph since in the backward pass I just want to compute the gradients of <code>loss</code> -&gt; <code>b2</code> -&gt; <code>b1</code>. Is there any ways to avoid <strong>Autograd</strong> and do the back-propagation manually?</p>
<p>Best!</p>

username: albanD
cooked: <ol>
<li>no if you chain them, then <code>func1</code> will also need to be differentiable to get gradients for module1.</li>
<li>You can use the information in <a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd" rel="nofollow noopener">this part of the doc</a> to create your own autograd Function. In your case, you can use <code>func1</code> and <code>func2</code> for the forwards and you will need to implement the backward for them.</li>
</ol>

username: vuhoangminh
cooked: <p>Thank you so much <a class="mention" href="/u/alband">@albanD</a></p>