Error on loss.backward() with RNN

username: zStupan
cooked: <p>This might be really dumb and the solution is probably obvious but I just can’t find it. I’m trying to train a char-rnn model but I’m getting a RuntimeError in the loss.backward() call:</p>
<p><code>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.</code></p>
<p>My model looks like this:</p>
<pre><code class="lang-python">class CharRNN(nn.Module):
  def __init__(self, vocab_size, embedding_dim=256, rnn_units=512, num_layers=3, dropout=0):
    super().__init__()
    self.embedding = nn.Embedding(vocab_size, embedding_dim)
    self.rnn = nn.GRU(embedding_dim, rnn_units, num_layers, dropout=dropout, batch_first=True)
    self.decoder = nn.Linear(rnn_units, vocab_size)
  
  def forward(self, inputs, state=None):
    x = self.embedding(inputs)
    output, state = self.rnn(x, state)
    output = self.decoder(output)
    return output, state
</code></pre>
<p>Custom loss class (I don’t know if this is correct, it’s supposed to be CrossEntropyLoss for inputs of shape <code>(batch_size, sequence_length, vocab_size)</code>):</p>
<pre><code class="lang-python">class CategoricalCrossEntropyLoss(nn.Module):
  def __init__(self):
    super().__init__()
    self.log_softmax = nn.LogSoftmax(-1)
  
  def forward(self, inputs, targets):
    log_probs = self.log_softmax(inputs)
    log_probs = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    return -log_probs.mean()
</code></pre>
<p>And my training loop looks like this:</p>
<pre><code class="lang-python">device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = CharRNN(vocab_size)
model.to(device)
loss_fn = CategoricalCrossEntropyLoss()
optimizer = torch.optim.RMSprop(model.parameters(), lr=0.002, alpha=0.95)
scheduler = StepLR(optimizer, 5, 0.95)

for epoch in range(50):
  running_loss = 0.0
  state = None
  model.train(True)
  for data in train_loader:
    inputs, targets = data[0].to(device), data[1].to(device)
    optimizer.zero_grad()
    output, state = model(inputs, state)
    loss = loss_fn(output, targets)
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
  scheduler.step()
  running_loss /= len(train_loader)
  
  # gradient clipping
  nn.utils.clip_grad.clip_grad_value_(model.parameters(), 5)

  # validate
  model.train(False)
  running_vloss = 0.0
  for data in val_loader:
    inputs, targets = data[0].to(device), data[1].to(device)
    output, state = model(inputs, state)
    vloss = loss_fn(output, targets)
    running_vloss += vloss.item()
    
  running_vloss /= len(val_loader)

  print(f'Epoch {epoch + 1}/{epochs}: Training Loss = {running_loss}; Validation Loss = {running_vloss};')
</code></pre>

username: Matteo_Ciotola
cooked: <p>I think the problem is that you clip the gradient after making the step with the optimizer.<br>
The correct may be:</p>
<aside class="quote no-group quote-modified" data-username="zStupan" data-post="1" data-topic="151209">
<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/zstupan/48/49416_2.png" class="avatar"> zStupan:</div>
<blockquote>
<pre><code class="lang-auto">for epoch in range(50):
  running_loss = 0.0
  state = None
  model.train(True)
  for data in train_loader:
    inputs, targets = data[0].to(device), data[1].to(device)
    optimizer.zero_grad()
    output, state = model(inputs, state)
    loss = loss_fn(output, targets)
    loss.backward()
    nn.utils.clip_grad.clip_grad_value_(model.parameters(), 5)
    optimizer.step()
</code></pre>
</blockquote>
</aside>
<p>You may try, it is a hypothesis.</p>

username: zStupan
cooked: <p>I had to detach the states of the rnn each iteration:</p>
<pre><code class="lang-python">if isinstance(state, torch.Tensor):
    state = state.detach()
else:  # for LSTM
    state = tuple(s.detach() for s in state)
</code></pre>
<p>But the gradient clipping was also in the wrong place. It should be before the step, as <a class="mention" href="/u/matteo_ciotola">@Matteo_Ciotola</a> pointed out.</p>