Custom layer giving problems when updating class parameter

William_Caster: <p>Hi everyone. I’m trying to implement a custom layer like this:</p>
<pre><code class="lang-auto">import torch.nn as nn
import math
import torch

class LinLayer(nn.Module):
    """ Custom Linear layer but mimics a standard linear layer """
    def __init__(self, size_in, size_out):
        super().__init__()
        self.size_in, self.size_out = size_in, size_out
        weights = torch.Tensor(size_out, size_in)
        self.weights = nn.Parameter(weights)  # nn.Parameter is a Tensor that's a module parameter.
        bias = torch.Tensor(size_out)
        self.bias = nn.Parameter(bias)

        self.hebb = torch.zeros((size_out, size_in), requires_grad=False)

        # initialize weights and biases
        nn.init.kaiming_uniform_(self.weights, a=math.sqrt(5)) # weight init
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weights)
        bound = 1 / math.sqrt(fan_in)
        nn.init.uniform_(self.bias, -bound, bound)  # bias init

    def forward(self, x):
        hebb = self.hebb

        w_times_x = torch.mm(x, self.weights.t() + hebb.t())

        yout = torch.add(w_times_x, self.bias)  # w times x + b

        self.hebb = torch.matmul(yout.t(), x)

        return yout
</code></pre>
<p>the problem is that when I try to update self.hebb = torch.matmul(yout.t(), x) at the end I get a “RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed).”. I have already tried detaching the self.hebb tensor but it make my loss a nan. How can I get this code to work?</p>

ptrblck: <aside class="quote no-group" data-username="William_Caster" data-post="1" data-topic="183649">
<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/william_caster/48/61958_2.png" class="avatar"> William_Caster:</div>
<blockquote>
<p>I have already tried detaching the self.hebb tensor but it make my loss a nan.</p>
</blockquote>
</aside>
<p>This sounds like the right approach since you are currently storing the computation graph in <code>self.hebb</code>. Since you are seeing a NaN loss after this fix you should narrow down where this invalid values comes from (e.g. in the forward pass, the gradients, or somewhere else).</p>

William_Caster: <p>Thx <a class="mention" href="/u/ptrblck">@ptrblck</a>! Indeed, the problem was coming from somewhere else (self.hebb was simply growing too large and generating nans at some point).</p>
<p>I would have a further question tho: in this framework of differential plasticity only these <strong>self.alpha</strong> and <strong>self.weight</strong> are optimized via backpropagation (so there’s no partial derivative of the loss wrt to self.hebb). Do I understand correctly that calling <strong>self.hebb.detach_()</strong> means that the partial derivative of the loss wrt <strong>self.hebb</strong> won’t be computed, but it won’t remove the effect of <strong>self.hebb</strong> when applying the chain rule to update <strong>self.weight</strong>  and <strong>self.alpha</strong>?</p>
<p>This is the forward function (<strong>self.alphas</strong> is one of the parameters of the model):</p>
<pre><code class="lang-auto">def forward(self, x):

        torch.clamp(self.alphas, min = -1.0, max = 1.0)

        w_times_x = torch.mm(x, self.weights.t() + (torch.mul(self.alphas, self.hebb)).t())

        yout = torch.add(w_times_x, self.bias)

        self.hebb = 0.1 * torch.matmul(yout.t(), x)

        self.hebb.detach_()

        return yout
</code></pre>

ptrblck: <p>Detaching <code>self.hebb</code> will avoid creating an increasing computation graph since otherwise you would backpropagate through the forward pass of the current iteration:</p>
<pre><code class="lang-python">        hebb = self.hebb
        w_times_x = torch.mm(x, self.weights.t() + hebb.t())
        yout = torch.add(w_times_x, self.bias)  # w times x + b
        self.hebb = torch.matmul(yout.t(), x)
</code></pre>
<p>as well as the previous one(s).</p>