Have to set "retain_graph = True" in .backward() to fit my model

Torch_GGG: <p>Hello. I’m building a model and try to fit with autograd. It’s an auto-regressive HMM, I have only 1 loss function, and I called optimizer.zero_grad() before I do loss.backward() and still getting the error</p>
<blockquote>
<p>RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</p>
</blockquote>
<p>It could only run with retain_graph set to True. It’s taking up a lot of RAM. Since I only have one loss, I think I don’t need it.</p>
<pre><code class="lang-auto">class ARHMM:

    def __init__(self, N, lags):

        self._N = torch.tensor(N)
        self._lags = torch.from_numpy(lags)
        self._max_lag = self._lags.max()

    def initialize(self):
        
        self._initial_matrix_raw = torch.rand((1, self._N), requires_grad= True)
        self._transition_matrix_raw = torch.rand((self._N, self._N), requires_grad= True)
        self._emission_parameter = torch.randn((self._N, len(self._lags) + 1), requires_grad = True)
        self._emission_sigma_raw = torch.randn((1, self._N), requires_grad = True)

        self._log_alpha = torch.zeros((self._T, self._N))

    def update_matrix(self):

        self._initial_matrix = self._initial_matrix_raw.clone().abs() / self._initial_matrix_raw.clone().abs().sum(dim = 1, keepdim = True)
        self._transition_matrix = self._transition_matrix_raw.clone().abs() / self._transition_matrix_raw.clone().abs().sum(dim = 1, keepdim = True)
        self._emission_sigma = self._emission_sigma_raw.clone().abs()

    def emission_logp(self, state, t):

        mu = self._emission_parameter[state, 0].clone()

        mu += torch.dot(self._y[t - self._lags], self._emission_parameter[state, 1:].to(torch.float64))

        x = self._y[t]

        sigma = self._emission_sigma[0, state].clone()

        logp = torch.tensor(-0.5) * (torch.tensor(2) * torch.tensor(np.pi) * sigma.square()).log() - (x - mu).square() / (torch.tensor(2) * sigma.square())

        return logp
    
    def log_sum_exp(self, acc):

        temp_max = acc.clone().max()

        return (acc.clone() - temp_max).exp().sum().log() + temp_max
        
    def calc_log_alpha(self):

        for i in torch.arange(self._N):

            self._log_alpha[self._max_lag, i] = self._initial_matrix[0, i].log() + self.emission_logp(i, self._max_lag)

        for t in torch.arange(self._max_lag + 1, self._T):

          for j in torch.arange(self._N):

            acc = self._log_alpha[t - 1] + self._transition_matrix[:, j].log() + self.emission_logp(j, t)

            self._log_alpha[t, j] = self.log_sum_exp(acc)

        self._negative_logp = torch.tensor(-1) * self.log_sum_exp(self._log_alpha[-1,:])

        print(self._negative_logp)
        print(self._emission_parameter)


    def fit(self, Y, iter = 500, lr = 0.001):

        self._y = torch.from_numpy(Y)
        self._T = self._y.shape[0]

        self.initialize()

        optimizer = torch.optim.Adam
        optimizer = optimizer([self._initial_matrix_raw,
                               self._transition_matrix_raw,
                               self._emission_parameter,
                               self._emission_sigma_raw],
                              lr=lr)

        for i in range(iter):

          self.update_matrix()
          self.calc_log_alpha()

          optimizer.zero_grad()
          self._negative_logp.backward()
          optimizer.step()

model = ARHMM(3, np.array([1,2,3]))
model.fit(np.array(obs), iter = 20)
model._emission_parameter
</code></pre>

ptrblck: <aside class="quote no-group" data-username="Torch_GGG" data-post="1" data-topic="174478">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/34f0e0/48.png" class="avatar"> Torch_GGG:</div>
<blockquote>
<p>and I called optimizer.zero_grad() before I do loss.backward() and still getting the error</p>
</blockquote>
</aside>
<p>Calling <code>zero_grad</code> will not change anything as it’s only deleting the gradients in newer PyTorch versions and setting them to zero in older ones.<br>
The error is raised by the actual <code>backward</code> call, which deletes the intermediate forward activations, which are needed for the gradient computation.<br>
In a common use case you would call <code>backward</code> only once per iteration after the forward pass and the loss computation. However, in your training code it seems you are trying to call <code>backward</code> multiple tmes after a single forward pass.<br>
In this case, you would need to keep the intermediates alive by using <code>retain_graph=True</code> and could call <code>backward</code> with <code>retain_graph=False</code> (default value) in the last step to clear the computation graph with its intermediates.</p>

Torch_GGG: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for your feedback. But I’m actually trying to call the backward() function once after the a single forward pass. The forward pass of my model is calc_log_alpha(). How should I modify the code to make it usable with retain_graph = False?</p>

ptrblck: <p>I guess <code>calc_log_alpha</code> is reusing tensors from previous iterations without detaching them from the old computation graph, which might cause the issue.<br>
You would need to check that all inputs (not parameters) are either newly created or detached from the previous computation graph in <code>calc_log_alpha</code>.</p>