How to circumvent need for retain_graph=True?

igygi: <p>Hi,</p>
<p>I am trying to train a GRU-based model. I encounter the following error when calling <code>loss.backward()</code>:<br>
<code>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></p>
<p>Setting <code>retain_graph=True</code> allows the model to run without any errors for a few iterations. However, the memory just keeps on going up until such point that I reach the <code>out of memory</code> error.</p>
<p>Here is my network (adopted from How to Train your Deep Multi-Object Tracker by Xu et al.):</p>
<pre><code class="lang-auto">class Munkrs(nn.Module):
    def __init__(self, element_dim, hidden_dim, target_size, biDirenction, minibatch, is_cuda, is_train=True):
        super(Munkrs, self).__init__()
        self.hidden_dim = hidden_dim
        self.bidirect = biDirenction
        self.minibatch = minibatch
        self.is_cuda = is_cuda

        # The LSTM takes word embeddings as inputs, and outputs hidden states
        # with dimensionality hidden_dim.
        self.lstm_row = nn.GRU(element_dim, hidden_dim, bidirectional=biDirenction, num_layers=2)
        self.lstm_col = nn.GRU(hidden_dim*2, hidden_dim, bidirectional=biDirenction, num_layers=2)

        # The linear layer that maps from hidden state space to tag space
        if biDirenction:
            # *2 directions * 2 ways concat
            self.hidden2tag_1 = nn.Linear(hidden_dim * 2, hidden_dim)
            self.hidden2tag_2 = nn.Linear(hidden_dim, 64)
            print('target_size', target_size)
            self.hidden2tag_3 = nn.Linear(64, target_size)
        else:
            # * 2 ways concat
            self.hidden2tag_1 = nn.Linear(hidden_dim, target_size)

        self.hidden_row = self.init_hidden(1)
        self.hidden_col = self.init_hidden(1)

    def forward(self, Dt):

        # Dt is of shape [batch, h, w]
        # input_row is of shape [h*w, batch, 1], [time steps, mini batch, element dimension]
        # row lstm #

        input_row = Dt.view(Dt.size(0), -1, 1).permute(1, 0, 2).contiguous()
        lstm_R_out, self.hidden_row = self.lstm_row(input_row, self.hidden_row)

        # column lstm #
        # lstm_R_out is of shape [seq_len=h*w, batch, hidden_size * num_directions]

        # [h * w*batch, hidden_size * num_directions]
        lstm_R_out = lstm_R_out.view(-1, lstm_R_out.size(2))

        # [h * w*batch, 1]
        # lstm_R_out = self.hidden2tag_1(lstm_R_out).view(-1, Dt.size(0))

        # [h,  w, batch, hidden_size * num_directions]
        lstm_R_out = lstm_R_out.view(Dt.size(1), Dt.size(2), Dt.size(0), -1)

        # col wise vector
        # [w,  h, batch, hidden_size * num_directions]
        input_col = lstm_R_out.permute(1, 0, 2, 3).contiguous()
        # [w*h, batch, hidden_size * num_directions]
        input_col = input_col.view(-1, input_col.size(2), input_col.size(3)).contiguous()
        lstm_C_out, self.hidden_col = self.lstm_col(input_col, self.hidden_col)

        # undo col wise vector
        # lstm_out is of shape [seq_len=time steps=w*h, batch, hidden_size * num_directions]

        # [h, w, batch, hidden_size * num_directions]
        lstm_C_out = lstm_C_out.view(Dt.size(2), Dt.size(1), Dt.size(0), -1).permute(1, 0, 2, 3).contiguous()

        # [h*w*batch, hidden_size * num_directions]
        lstm_C_out = lstm_C_out.view(-1, lstm_C_out.size(3))

        # [h*w, batch, 1]
        tag_space = self.hidden2tag_1(lstm_C_out)
        tag_space = self.hidden2tag_2(tag_space)
        tag_space = self.hidden2tag_3(tag_space).view(-1, Dt.size(0))
        tag_scores = torch.sigmoid(tag_space)
        # tag_scores is of shape [batch, h, w] as Dt

        return_this = tag_scores.view(Dt.size(1), Dt.size(2), -1).permute(2, 0, 1).contiguous()
        
        return return_this

</code></pre>
<p>And here is my training loop:</p>
<pre><code class="lang-auto">for epoch in range(num_epochs):
	print('Epoch {}/{}'.format(epoch, num_epochs-1))
	print('-' * 10)

	for phase in ['train', 'val']:
		if phase == 'train':
			model.train()
		else:
			model.eval()

		running_loss = 0.0

		for i, data in enumerate(dataloader[phase]):
			model.detach_()
			print('i', i)

			# get the inputs 
			x = data['input'].to(device)
			y_gt = data['label'].to(device)

			# zero the parameter gradients
			optimizer.zero_grad()

			# forward, track history if only in train
			with torch.set_grad_enabled(phase == 'train'):
				y_pred = model(x)
				#print('input shape', x.shape)
				y_pred = y_pred.view(batch_size, x.shape[1]*x.shape[2])
				y_gt = y_gt.view(batch_size, x.shape[1]*x.shape[2])
				loss = focalLoss(y_pred, y_gt)

				if phase == 'train':
					loss.backward()
					optimizer.step()

				with open(running_loss_log, 'a') as f:
					print(running_loss)
					df = pd.DataFrame([[phase, epoch, i, loss]], columns=['split', 'epoch', 'i', 'loss'])
					df.to_csv(f, header=False)

			# print statistics
			with torch.no_grad():
				running_loss += loss

			if i % 10 == 9:
				print('[%d, %5d] Running loss: %.3f' % (epoch+1, i+1, running_loss/i))

		if phase == 'train':
			exp_lr_scheduler.step()

		epoch_loss = running_loss / dataset_sizes[phase]
		print('{} Epoch Loss: {:.4f}'.format(phase, epoch_loss))

		with open(epoch_loss_log, 'a') as f:
			df = pd.DataFrame([[phase, epoch, epoch_loss]], columns=['split', 'epoch', 'loss'])
			df.to_csv(f, header=False)

		# deep copy the model according to loss
		if phase == 'val' and epoch_loss &lt; best_loss:
			best_loss = epoch_loss
			best_model_wts = copy.deepcopy(model.state_dict())

	if epoch % 2 == 0:
		latest_model = copy.deepcopy(model.state_dict())

		torch.save(best_model_wts, to_save + '_bestloss.pth')
		torch.save(latest_model, to_save + 'latest.pth')


print('Finished Training')
</code></pre>