Custom backward pass

sadegh: <p>Hi everyone,</p>
<p>I have a neural network that contains some complex operations (e.g., f3 and f6 in this figure). In the forward pass, we use all the black arrows (use all f1, f2, …, f7), however, in the backward pass, instead of using f3 and f6, we want to use g3 and g6, which are pretty close f3 and f6 but are simpler.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a2acbb484376eceb75809be0d14130a919df582c.png" data-download-href="https://discuss.pytorch.org/uploads/default/a2acbb484376eceb75809be0d14130a919df582c" title="Untitled Diagram"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a2acbb484376eceb75809be0d14130a919df582c.png" alt="Untitled Diagram" data-base62-sha1="nd5qp9lEdiu8RNgejPmrVpD2D2k" width="690" height="133" data-dominant-color="F4F5F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Untitled Diagram</span><span class="informations">891×172 5.89 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Both f3 and f6 contain multiple layers and operations, so I define them in another class:</p>
<pre><code class="lang-auto">class CustomLayer_F(nn.Module):
    def __init__(self):
        super(CustomLayer_F, self).__init__()
           # defining all variables and operations
        
    def forward(self, input):
        # computing the forward pass
</code></pre>
<p>And for g3 and g6, I defnie:</p>
<pre><code class="lang-auto">class CustomLayer_G(nn.Module):
    def __init__(self):
        super(CustomLayer_G, self).__init__()
           # defining all variables and operations, different from F
        
    def forward(self, input):
        # computing the forward pass, different from F
</code></pre>
<p>In order to use <code>CustomLayer_F</code> in the forward pass and <code>CustomLayer_G</code>’s gradients in the backward pass, I want something like:</p>
<pre><code class="lang-auto">class CustomForwardBackward(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
          # do whatever we do in CustomLayer_F

    @staticmethod
    def backward(ctx, grad_output):
        # forward through CustomLayer_G and use the gradients of G instead of F
        # However, I don't know how to use grad_output here!!!
</code></pre>
<p>Does anyone have any clue how can I use <code>CustomForwardBackward</code> to forward like CustomLayer_F but backward through CustomLayer_G and use its gradients?</p>
<p>Thank you very much!</p>

ptrblck: <p>Cross-posting from Twitter:<br>
Check out <a href="https://discuss.pytorch.org/t/relu-with-leaky-derivative/32818/2">this post</a> and see, if it’s applicable for this use case.</p>

sadegh: <p>Thanks for the link and the discussion on twitter! It was actually helpful, however, a simpler solution that worked for me was this:</p>
<pre><code class="lang-auto">class CustomForwardBackward(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
          with torch.enable_grad():
              output = ctx.g3.forward(input)
              ctx.save_for_backward(input, output)
          return ctx.f3.forward(input)
    @staticmethod
    def backward(ctx, grad_output):
          input, output = ctx.saved_tensors
          output.backward(grad_output, retain_graph=True)
          return input.grad
</code></pre>