Custom autograd.Function: backward pass not called

Technics: <p>Good afternoon!</p>
<p>I’ve had this problem in my other thread already, but it isn’t really related, so I moved it to a new thread.</p>
<p>Whenever I try using a custom autograd.Function (i.e., a function that has an explicit backward pass defined), and I combine it with any torch.nn module, the backward pass is never properly executed.<br>
Here’s an MWE containing a simple identity transform as an example:</p>
<hr>
<pre><code class="lang-auto">import torch
import torch.autograd
import torch.nn



# custom autograd function
class Identity(torch.autograd.Function):
    def __init__(self):
        super(Identity, self).__init__()

    @staticmethod
    def forward(self, input):
        return input.clone()        # REASON FOR ERROR: forgot to .clone() here
    
    @staticmethod
    def backward(self, grad_output):
        print("backward")
        grad_input = grad_output.clone()
        return grad_input

identity = Identity.apply





input = torch.autograd.Variable(torch.randn(1,2,20,20).float(), requires_grad=True)
target = torch.autograd.Variable(torch.randn(1,2,20,20).float(), requires_grad=False)


# forward pass:
# - if any torch.nn function is called before the custom Function,
#   its backward pass is never executed.


intermediate = input                           # calling backward upon this works fine
# intermediate = torch.nn.ReLU()(input)        # uncomment this line and it won't work anymore


pred = identity(intermediate)


loss = torch.nn.MSELoss()(pred,target)
loss.backward()                                # this should always print "backward" to the command line, but it doesn't...</code></pre>
<hr>
<p>I am running the latest version of PyTorch (Anaconda says 0.2.0, py27hc03bea1_4cu80  [cuda80]  soumith), with CUDA 8.0 and Python 2.7.</p>
<p>With some PyTorch versions this seemed to have worked, but now it won’t anymore.</p>
<p>Anybody knows what’s the matter, or else what I’m doing wrong?</p>
<p>Thank you very much for any answers!</p>
<p><strong>EDIT:</strong> Silly me. I’ve been trying to solve this for quite some time, and literally 5 minutes after opening this thread I’ve discovered my mistake: I forgot to clone the inputs during the forward pass (code changed above), now it works.<br>
Sorry for the spamming; thread can be closed.</p>

albanD: <p>If you don’t want to put the <code>clone</code>, you should mark that the input and output share some elements with the <code>mark_shared_storage</code> function like <a href="https://github.com/pytorch/pytorch/blob/9735ddd89904be67e4924f6e69ac6df4cf40c083/torch/autograd/_functions/tensor.py#L93">here</a>.</p>

Akash_Goel: <p>Hey could you explain me why i need to return input.clone()</p>

Technics: <p><a class="mention" href="/u/alband">@albanD</a>: thanks a lot for the additional input! I’m quite fine with cloning, for my variables are not immensely huge.</p>
<p><a class="mention" href="/u/akash_goel">@Akash_Goel</a>: I’m obviously not an expert, but I assume the reason is the directed graph that is constructed behind the scene while running the forward pass on a variable. Every variable is linked to a predecessor and a successor. During the backward pass, this directed graph is traversed in inverted order.<br>
If one doesn’t clone the variable, its executor function is probably just overwritten, which leads to the skipping of one of the building blocks. By cloning, however, a new variable is appended that is assigned to the custom block, and this preserves the execution order.<br>
Correct me if I’m wrong.</p>

Akash_Goel: <p><a class="mention" href="/u/technics">@Technics</a>: I think you are correct since there is no operation performing on “input” no new operation has been added to the graph and since operation is not in the graph the backward function won’t be called by the graph during backpropogate <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Chun_Li: <p>My question in post</p>
<aside class="quote quote-modified" data-post="1" data-topic="9129">
  <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/chun_li/48/55_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/get-error-message-maskedfill-cant-differentiate-the-mask/9129">Get error message "MaskedFill can't differentiate the mask"</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    My a program was able to run in earlier version 0.2.0+0eec332 of pytorch (date back to Oct. 6, 2017). 
Yesterday and today I upgraded my pytorch to the latest version from source, and now I can’t run 
the program. The error message I got is as follows 
File “/usr/local/lib/python3.5/dist-packages/torch/autograd/variable.py”, line 85, in setitem 
return MaskedFill.apply(self, key, value, True) 
File “/usr/local/lib/python3.5/dist-packages/torch/autograd/_functions/tensor.py”, line 483, in forward …
  </blockquote>
</aside>

<p>also related to custom autograd.Function and backward pass. Inspired by your code,<br>
I modified my code and make it works now in the newer version of pytorch. However, I do not<br>
clone the input during the forward pass, and it still works. Interesting, but also confusing me!</p>

kevinkepp: <p>Had a similar problem. For me, <code>mark_dirty</code> worked perfectly. Makes <code>backward</code> getting called without needing to clone the input tensor or store pairs of tensors with <code>mark_shared_storage</code>.</p>
<pre><code class="lang-auto">class ManipulateGradient(Function):
	def forward(self, input):
		self.mark_dirty(input)
		return input

	def backward(self, grad_out):
		# manipulate gradient here
		return grad_out + 0.42</code></pre>

Adam_Dziedzic: <p>Does it really work for you?</p>
<pre><code class="lang-auto">from torch.autograd import Function
from torch import tensor

class ManipulateGradient(Function):

    def forward(self, input):
        self.mark_dirty(input)
        print("forward Function")
        return input

    def backward(grad_out):
        # manipulate gradient here
        print("backward Function")
        return grad_out + 0.42

if __name__ == "__main__":
    layer = ManipulateGradientModule()
    input = tensor([1.0, 2.0, 3.0], requires_grad=True)
    result = layer.forward(input)
    print("result: ", result)
    dout = tensor([0.1, -0.1, 0.2])
    result.backward(dout)
    print("result.grad: ", input.grad)
</code></pre>
<p>Gives me:</p>
<pre><code class="lang-auto">forward Function
Traceback (most recent call last):
    result = layer.apply(input)
RuntimeError: a leaf Variable that requires grad has been used in an in-place operation.
</code></pre>
<p>Okay, let’s remove</p>
<pre><code class="lang-auto">self.mark_dirty(input)
</code></pre>
<p>Then we get:</p>
<pre><code class="lang-auto">forward Function
result:  tensor([1., 2., 3.], grad_fn=&lt;FunctionBackward&gt;)
Traceback (most recent call last):
  File "manipulate_gradient_non_static.py", line 22, in &lt;module&gt;
    result.backward(dout)
  File "python3.6/site-packages/torch/tensor.py", line 102, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "python3.6/site-packages/torch/autograd/__init__.py", line 90, in backward
    allow_unreachable=True)  # allow_unreachable flag
  File "python3.6/site-packages/torch/autograd/function.py", line 76, in apply
    return self._forward_cls.backward(self, *args)
  File "python3.6/site-packages/torch/autograd/function.py", line 180, in backward
    raise NotImplementedError
NotImplementedError

</code></pre>
<p>Let’s write a proper Function with <span class="mention">@staticmethods</span>:</p>
<pre><code class="lang-auto">class ManipulateGradientStatic(Function):

    @staticmethod
    def forward(ctx, input):
        # ManipulateGradientStatic.mark_dirty(input)
        print("forward static")
        return input

    @staticmethod
    def backward(ctx, grad_out):
        # manipulate gradient here
        print("backward static")
        return grad_out + 0.42

class ManipulateGradientModule(torch.nn.modules.Module):
    def __init__(self):
        super(ManipulateGradientModule, self).__init__()

    def forward(self, input):
        print("forward Module")
        return ManipulateGradientStatic.apply(input)

if __name__ == "__main__":
    layer = ManipulateGradientModule()
    input = tensor([1.0, 2.0, 3.0], requires_grad=True)
    result = layer.forward(input)
    print("result: ", result)
    dout = tensor([0.1, -0.1, 0.2])
    result.backward(dout)
    print("result.grad: ", input.grad)
</code></pre>
<p>Gives:</p>
<pre><code class="lang-auto">forward Module
forward static
result:  tensor([1., 2., 3.], grad_fn=&lt;ManipulateGradientStaticBackward&gt;)
backward static
result.grad:  tensor([0.5200, 0.3200, 0.6200])
</code></pre>

Adam_Dziedzic: <p>Is there any way to debug why my statcic backward method (much more complicated than the example above) is not called during backprop?</p>

colesbury: <p>You can look at the chain of <code>grad_fn</code> and <code>next_functions</code>:</p>
<pre><code class="lang-auto">print(result.grad_fn)
print(result.grad_fn.next_functions)
print(result.grad_fn.next_functions[0][0].next_functions)
...
</code></pre>
<p>Usually the issue has something to do with either a <code>detach()</code> call (which stops backpropagation) or returning the input unmodified in <code>forward</code>. Try returning <code>input.clone()</code> to see if it fixes the issue.</p>

janfreyberg: <p>Hi, apologies for reviving an old thread, but even with the solutions indicated here I am still a bit confused. I have implemented a <a href="https://github.com/janfreyberg/superintendent" rel="noopener nofollow ugc">gradient reversal layer</a>. I started by implementing a function:</p>
<pre data-code-wrap="python"><code class="lang-python"># functional.py
from torch.autograd import Function


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

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = None
        if ctx.needs_input_grad[0]:
            grad_input = -grad_output.clone()
        return grad_input


revgrad = RevGrad.apply
</code></pre>
<p>I then use this function in a layer:</p>
<pre data-code-wrap="python"><code class="lang-python"># __init__.py
from .functional import revgrad
from torch.nn import Module


class RevGrad(Module):
    def __init__(self, *args, **kwargs):
        """
        A gradient reversal layer.

        This layer has no parameters, and simply reverses the gradient
        in the backward pass.
        """

        super().__init__(*args, **kwargs)

    def forward(self, input_):
        return revgrad(input_)

</code></pre>
<p>This all works as expected, and the following test (run with pytest) succeeds:</p>
<pre data-code-wrap="python"><code class="lang-python"># test_layer.py

import copy
import torch
from pytorch_revgrad import RevGrad


def test_gradients_inverted():
    network = torch.nn.Sequential(torch.nn.Linear(5, 3), torch.nn.Linear(3, 1))
    revnetwork = torch.nn.Sequential(copy.deepcopy(network), RevGrad())

    inp = torch.randn(8, 5)
    outp = torch.randn(8)

    criterion = torch.nn.MSELoss()
    criterion(network(inp), outp).backward()
    criterion(revnetwork(inp), outp).backward()
    assert all(
        (p1.grad == -p2.grad).all()
        for p1, p2 in zip(network.parameters(), revnetwork.parameters())
    )
</code></pre>
<p>So far so good. However, when I run this with the pytest-cov plugin, to check the coverage of my source code, it makes it look as if the backward call in my function is never called:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/b/ba876a411b614da31c7860e9867ba9ba8d6597e3.png" alt="image" data-base62-sha1="qC6VXQij6iigrDzLlHf4ixZ8Hu3" width="601" height="400"></p>
<p>I’m very confused by this - the gradients do get reversed, as is clear from the test succeeding. But somehow this happens without the backward call being made? Am I missing something fundamental here?</p>

albanD: <p>Hi,</p>
<p>I’m not sure how the coverage is checked, but if it checks if a python function calls the backward it will never happen. It is called directly by some cpp code. So the coverage tool might miss it.</p>

janfreyberg: <p>Thanks, that’s extremely helpful to know!</p>
<p>I think that is the issue, as far as I know coverage.py uses <a href="https://docs.python.org/3/library/sys.html#sys.settrace" rel="nofollow noopener">sys.settrace</a> to do this, so that explains things.</p>

deathcrush: <p>Sorry to ask a silly question, but why is it necessary to have this layer in the first place. Could you not negate the loss function directly? Then the gradients would be flipped automatically, with no extra lines of code?</p>

wandering007: <p>Hi, I am also faced with the similar issue. Have you got any solutions? Thanks.</p>

janfreyberg: <p>Hi, the answer is up there - it’s being called, it’s just hidden from the python code because it’s called from C.</p>

janfreyberg: <p>Sorry, only just saw this.</p>
<p>It’s for when you want to negate the gradients selectively for some parts of the network; multiplying the loss by -1 works but doesn’t give you any choice over which parts of the network get updated with the inverted gradients</p>