Gradient scaling, reversal

AlbertZeyer: <p>I wonder about the best way how to implement gradient reversal, or in general gradient scaling (reversal is the special case of using factor -1 then).</p>
<p>Related:</p><aside class="quote quote-modified" data-post="1" data-topic="3589">
  <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/bartolomed/48/3365_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/solved-reverse-gradients-in-backward-pass/3589">[Solved] Reverse gradients in backward pass</a> 
  </div>
  <blockquote>
    Hello everyone, 
I am working on building a DANN (Ganin et al. 2016) in PyTorch. This model is used for domain adaptation, and forces a classifier to only learn features that exist in two different domains, for the purpose of generalization across these domains. The DANN uses a Gradient Reversal layer to achieve this. 
I have seen some suggestions on this forum on how to modify gradients manually. However, I found it difficult to apply in my case, as the gradients are reversed midway through the…
  </blockquote>
</aside>
<aside class="quote" data-post="1" data-topic="134810">
  <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/hadaev8/48/16280_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-reverse-gradient-sign-during-backprop/134810">How to reverse gradient sign during backprop?</a> 
  </div>
  <blockquote>
    For example, I have models net0 and net1 working like this 
out0 = net0(inp0)
out1 = net0(net1(net0(inp1)))
loss = criterion(out0, out1, target)

I want to reverse the gradient sign from the output of model net1 but calculate net0(inp0) as usual. 
In simple case, I would do 
out0 = net0(inp0)
out1 = net0(net1(inp1))
loss = criterion(out0, out1, target)
loss.backward()
[p.grad.data.neg_() for p in net1.parameters()]
opt.step()

Can I do some hook to reverse gradients from output of net1?
  </blockquote>
</aside>

<p>Existing implementations:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/fairseq/blob/100cd91db19bb/fairseq/modules/grad_multiply.py">
  <header class="source">

      <a href="https://github.com/facebookresearch/fairseq/blob/100cd91db19bb/fairseq/modules/grad_multiply.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/fairseq/blob/100cd91db19bb/fairseq/modules/grad_multiply.py" target="_blank" rel="noopener nofollow ugc">facebookresearch/fairseq/blob/100cd91db19bb/fairseq/modules/grad_multiply.py</a></h4>


      <pre><code class="lang-py"># Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch


class GradMultiply(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale):
        ctx.scale = scale
        res = x.new(x)
        return res

    @staticmethod
    def backward(ctx, grad):
        return grad * ctx.scale, None
</code></pre>




  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubblob" data-onebox-src="https://github.com/janfreyberg/pytorch-revgrad/blob/449fa763a76d/src/pytorch_revgrad/functional.py">
  <header class="source">

      <a href="https://github.com/janfreyberg/pytorch-revgrad/blob/449fa763a76d/src/pytorch_revgrad/functional.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/janfreyberg/pytorch-revgrad/blob/449fa763a76d/src/pytorch_revgrad/functional.py" target="_blank" rel="noopener nofollow ugc">janfreyberg/pytorch-revgrad/blob/449fa763a76d/src/pytorch_revgrad/functional.py</a></h4>


      <pre><code class="lang-py">from torch.autograd import Function


class RevGrad(Function):
    @staticmethod
    def forward(ctx, input_, alpha_):
        ctx.save_for_backward(input_, alpha_)
        output = input_
        return output

    @staticmethod
    def backward(ctx, grad_output):  # pragma: no cover
        grad_input = None
        _, alpha_ = ctx.saved_tensors
        if ctx.needs_input_grad[0]:
            grad_input = -grad_output * alpha_
        return grad_input, None


revgrad = RevGrad.apply
</code></pre>




  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubblob" data-onebox-src="https://github.com/tadeephuy/GradientReversal/blob/5d9857d63/gradient_reversal/functional.py">
  <header class="source">

      <a href="https://github.com/tadeephuy/GradientReversal/blob/5d9857d63/gradient_reversal/functional.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/tadeephuy/GradientReversal/blob/5d9857d63/gradient_reversal/functional.py" target="_blank" rel="noopener nofollow ugc">tadeephuy/GradientReversal/blob/5d9857d63/gradient_reversal/functional.py</a></h4>


      <pre><code class="lang-py">from torch.autograd import Function

class GradientReversal(Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.save_for_backward(x, alpha)
        return x
    
    @staticmethod
    def backward(ctx, grad_output):
        grad_input = None
        _, alpha = ctx.saved_tensors
        if ctx.needs_input_grad[0]:
            grad_input = - alpha*grad_output
        return grad_input, None
revgrad = GradientReversal.apply</code></pre>




  </article>

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

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

<p>Some questions on this code:</p>
<p>Fairseq just does <code>ctx.scale = scale</code>, while the other implementations use <code>ctx.save_for_backward(input_, alpha_)</code>. What’s the difference? What is better?</p>
<p>Fairseq uses <code>res = x.new(x)</code> but the others do not. Why is this needed? What does it actually do? I did not found the documentation on <code>Tensor.new</code>.</p>
<p>The other implementations check for <code>ctx.needs_input_grad[0]</code> in the backward pass but Fairseq does not. Is this not needed?</p>