Saving gradients of operators

Joe_Harrison: <p>I’m trying to save the intermediate gradients of a graph as done in table 3 in <a href="https://www.jmlr.org/papers/volume18/17-468/17-468.pdf" rel="noopener nofollow ugc">Baydin et al.</a>. I have custom modules where I’d like to know the gradients. To do this I registered a backward hook on these modules. The gradients of the leaf-nodes are correct, but the gradients I get from the hook aren’t. How do I get the gradients that flow through operators?</p>
<p>Here’s what I have now:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Node(nn.Module):
    def __init__(self):
        super(Node, self).__init__()
        self.children = []
        self.name = 0
    
    def forward(self, x):
        pass
    
    def get_subtree(self):
        subtree = []
        self._get_subtree_recursive(subtree)
        return subtree
    
    def _get_subtree_recursive(self, subtree):
        subtree.append(self)
        for c in self.children:
            c._get_subtree_recursive(subtree)

class Plus(Node):
    def __init__(self, child0, child1):
        super(Plus, self).__init__()
        self.children = [child0, child1]
        self.name = "Plus"
    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        return c_outputs[0] + c_outputs[1]

class Minus(Node):
    def __init__(self, child0, child1):
        super(Minus, self).__init__()
        self.children = [child0, child1]
        self.name = "Minus"
    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        return c_outputs[0] - c_outputs[1]

class Sin(Node):
    def __init__(self, child0):
        super(Sin, self).__init__()
        self.children = [child0]
        self.name = "Sin"
    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        return torch.sin(c_outputs[0])
    
class Ln(Node):
    def __init__(self, child0):
        super(Ln, self).__init__()
        self.children = [child0]
        self.name = "Ln"
    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        return torch.log(c_outputs[0])

class Multiply(Node):
    def __init__(self, child0, child1):
        super(Multiply, self).__init__()
        self.children = [child0, child1]
        self.name = "Multiply"
    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        return c_outputs[0] * c_outputs[1]

class Const(Node):
    def __init__(self, value):
        super(Const, self).__init__()
        self.value = nn.Parameter(torch.tensor([value], requires_grad=True))
        self.name = "Const({0:.3f})".format(self.value.item())
    
    def forward(self, x):
        return self.value.repeat(x.size(0),1)

# Dict to save grads in hook
grads = {}

def savegrad(self, grad_input, grad_output):
    if grad_input[0] is not None:
        grads[id(self)] = (grad_input[0], self.name)

# Expression with common constants
consts = [Const(2.), Const(5.)]
expression = Minus(Plus(Ln(consts[0]),Multiply(consts[0],consts[1])), Sin(consts[1]))
expression_nodes = expression.get_subtree()

# Random input
x = torch.tensor([[42.]])

for idx, c in enumerate(expression_nodes):
    c.register_backward_hook(savegrad)

z = expression(x)
z.backward()

# These values correspond with Baydin et al.
print(consts[0].value.grad)
print(consts[1].value.grad)

# These don't
for idx, c in enumerate(expression_nodes):
    if id(c) in grads.keys():
        print(grads[id(c)])
</code></pre>

KFrank: <p>Hi Joe!</p>
<aside class="quote no-group" data-username="Joe_Harrison" data-post="1" data-topic="150120" 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/joe_harrison/48/5534_2.png" class="avatar"> Joe_Harrison:</div>
<blockquote>
<p>How do I get the gradients that flow through operators?</p>
</blockquote>
</aside>
<p>I haven’t looked at your code, but could <a href="https://pytorch.org/docs/1.11/generated/torch.Tensor.retain_grad.html" rel="noopener nofollow ugc">retain_grad()</a> work for your<br>
use case?</p>
<p>A simple example:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.11.0'
&gt;&gt;&gt; v = torch.tensor ([1., 2., 3.], requires_grad = True)
&gt;&gt;&gt; t = torch.outer (v, v)
&gt;&gt;&gt; t.retain_grad()
&gt;&gt;&gt; t.sum().backward()
&gt;&gt;&gt; t.grad
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Joe_Harrison: <p>Thank you! Both work, but the reason why the values differ from the paper is that the last gradient overwrites .grad. Some nodes are used twice so for these nodes a history of gradients needs to be retained. The values now correspond with the paper.</p>
<p>Sample Node and hook</p>
<pre><code class="lang-auto">class Node(nn.Module):
    def __init__(self):
        super(Node, self).__init__()
        self.children = []
        self.name = 0
        self.ret_grad = False
        self.output = []
    
    def retain_grad(self, retain_grad=True):
        self.ret_grad = retain_grad
    
    def forward(self, x):
        pass
    
    def get_subtree(self):
        subtree = []
        self._get_subtree_recursive(subtree)
        return subtree
    
    def _get_subtree_recursive(self, subtree):
        subtree.append(self)
        for c in self.children:
            c._get_subtree_recursive(subtree)

class Plus(Node):
    def __init__(self, child0, child1):
        super(Plus, self).__init__()
        self.children = [child0, child1]
        self.name = "Plus"

    
    def forward(self, x):
        c_outputs = [c(x) for c in self.children]
        output = c_outputs[0] + c_outputs[1]
        if self.ret_grad:
            output.retain_grad()
            self.output.append(output)
        return output

def savegrad(self, grad_input, grad_output):
    if grad_input[0] is not None:
        if id(self) not in grads.keys():
            grads[id(self)] = []
        grads[id(self)].append(grad_input[0])
</code></pre>