Pytorch - lstm yields retain_graph error - how do I get around this?

Imran_Rashid1: <p>I am training a simple LSTM model however pytorch gives me an error saying that I need to set retain_graph=True. However this takes the model longer to train and I do not think I need to do this.</p>
<pre><code class="lang-auto">class SequenceModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size = 3, hidden_size = 3, bidirectional=False)
        self.hidden = (torch.randn(1, 1, 3).double(), torch.randn(1, 1, 3).double())

    def forward(self,x):
        lstm_out, self.hidden = self.lstm(x.view(-1, 1, 3),self.hidden)
        return lstm_out

    def loss(self,logits,labels):
        return F.cross_entropy(logits, labels)

model = SequenceModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

model = model.double()

model.train()
epochs = 1000
for epoch in tqdm(range(epochs)):
    optimizer.zero_grad()


    logits = model(inputs)
    logits = logits.reshape(-1,3)
    loss = model.loss(logits,outputs.long())

    loss.backward() 
    optimizer.step()
</code></pre>
<p>The error I get is:</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.
</code></pre>
<p>I do not want to use retain_graph=True because the training takes longer to run. I do not think that my simple LSTM should need the <code>retain_graph=True</code>. What am I doing wrong?</p>

albanD: <p>Hi,</p>
<p>The problem is that the hidden layers in your model are shared from one invocation to the next. And so they are all linked.<br>
In particular, because the LSTM module runs the whole forward, you do not need to save the final hidden states:</p>
<pre><code class="lang-auto">lstm_out, _ = self.lstm(x.view(-1, 1, 3),self.hidden)
</code></pre>

Imran_Rashid1: <p>Thanks, in that case should I remove <code>self.hidden</code> altogether, i.e.</p>
<p><code>lstm_out, _ = self.lstm(x.view(-1, 1, 3))</code></p>
<p>so that it doesn’t take the random <code>self.hidden</code> from the constructor?</p>

albanD: <p>I think usually the default of them being all 0 is good yes.<br>
So unless you want to do something fancy, you don’t need to provide them !</p>

Imran_Rashid1: <p>Sorry I meant that after every training step, the lstm should taken the most up-to-date hidden layers( h,c) right?</p>
<p>That’s why I was passing the up-to-date self.hidden as an argument for the lstm input. Or am I missing something?</p>

albanD: <p>In the most basic for of lstm, IIRC, the hidden layer value for the first iteration should be full of 0.<br>
Since the module you use performs all the iterations in one call. You only ever want to provide a Tensor full of 0s for it.</p>
<p>If you want to train the initial hidden layer value, you can declare them as nn.Parameter and given them as output to the lstm (but still ignore the returned value).</p>

Imran_Rashid1: <p>But after the first iteration/step of training should the first hidden layer still be all zeroes. I’m guessing not since when we’re training the first hidden layer weight values change. In that case:</p>
<p>On my first training step:</p>
<p>my h0 has changed from all zeroes to some non-zero weights - let’s say h0*.</p>
<p>On my second training step:</p>
<p>how do I specify that I want the updated h0*</p>
<p>i.e. lstm_out, _ = self.lstm(x.view(-1, 1, 3),h0*)</p>
<p>without sharing the hidden layers from one invocation the next.</p>
<p>I’m sure I’m confusing something here - but I just wanted to double check.</p>

albanD: <p>If you want to learn the hidden layer initial state, starting from 0:</p>
<pre><code class="lang-python">class SequenceModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size = 3, hidden_size = 3, bidirectional=False)
        self.hidden = nn.ParameterList((nn.Parameter(torch.zeros(1, 1, 3),
                                        nn.Parameter(torch.zeros(1, 1, 3)))

    def forward(self,x):
        lstm_out, _ = self.lstm(x.view(-1, 1, 3),self.hidden)
        return lstm_out
</code></pre>
<p>And now it will be picked up by the <code>.parameters()</code> and so the optimizer will take care of updating your hidden state with the rest of the parameters.</p>

Imran_Rashid1: <p>Thank you very much. I’m guessing this is the proper way to train an LSTM when you’re training with more than one step/iteration.</p>

albanD: <p>Not necessarily. Keeping them as <code>0</code> for the first step all the time is valid as well.<br>
And reduces the number of parameters so might help avoid overfitting.</p>

Imran_Rashid1: <p>Thanks. Learn something new everyday <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>