PyTorch：寻找已被就地操作修改的梯度计算所需的变量 - 多任务学习

Question

I recently did a massive refactor to my PyTorch LSTM code, in order to support multitask learning. I created an MTLWrapper , which holds a BaseModel (which can be one of several variations on a regular LSTM.network), which remained the same as it was before the refactor, minus a linear hidden2tag layer (takes hidden sequence and converts to tag space), which now sits in the wrapper. The reason for this is that for multitask learning, all the parameters are shared, except for the final linear layer, which I have one of for each task. These are stored in a nn.ModuleList , not just a regular python list.

What happens now is that my forward pass returns a list of tag scores tensors (one for each task), rather than a single tensor of the tag scores for a single task. I compute the losses for each of these tasks and then try to backpropagate with the average of these losses (technically also averaged over all the sentences of a batch, but this was true before the refactor too). I call model.zero_grad() before running the forward pass on each sentence in a batch.

I don't know exactly where it happened, but after this refactor, I started getting this error (on the second batch):

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.

Following the advice, I added the retain_graph=True flag, but now I get the following error instead (also on the second backward step):

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [100, 400]], which is output 0 of TBackward, is at version 2; expected version 1 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

The hint in the backtrace is not actually helpful, because I have no idea where a tensor of the shape [100, 400] even came from - I don't have any parameters of size 400. I have a sneaky suspicion that the problem is actually that I shouldn't need the retain_graph=True , but I have no way to confirm that vs. finding the mystery variable that is being changed according to the second error. Either way, I'm at a complete loss how to solve this issue. Any help is appreciated!

Code snippets:

import torch
import torch.nn as nn
import torch.nn.functional as F

class MTLWrapper(nn.Module):
    
    def __init__(self, embedding_dim, hidden_dim, dropout,..., directions=1, device='cpu', model_type):
        super(MTLWrapper, self).__init__()
        self.base_model = model_type(embedding_dim, hidden_dim, dropout, ..., directions, device)
        self.linear_taggers = []
        for tagset_size in tagset_sizes:
            self.linear_taggers.append(nn.Linear(hidden_dim*directions, tagset_size))
        self.linear_taggers = nn.ModuleList(self.linear_taggers)

    def init_hidden(self, hidden_dim):
        return self.base_model.init_hidden(hidden_dim)

    def forward(self, sentence):
        lstm_out = self.base_model.forward(sentence)
        tag_scores = []
        for linear_tagger in self.linear_taggers:
            tag_space = linear_tagger(lstm_out.view(len(sentence), -1))
            tag_scores.append(F.log_softmax(tag_space))
        tag_scores = torch.stack(tag_scores)
        return tag_scores

Inside the train function:

for i in range(math.ceil(len(train_sents)/batch_size)):
    batch = r[i*batch_size:(i+1)*batch_size]
    losses = []
    for j in batch:

        sentence = train_sents[j]
        tags = train_tags[j]

        # Step 1. Remember that Pytorch accumulates gradients.
        # We need to clear them out before each instance
        model.zero_grad()

        # Also, we need to clear out the hidden state of the LSTM,
        # detaching it from its history on the last instance.
        model.hidden = model.init_hidden(hidden_dim)

        sentence_in = sentence
        targets = tags

        # Step 3. Run our forward pass.
        tag_scores = model(sentence_in)

        loss = [loss_function(tag_scores[i], targets[i]) for i in range(len(tag_scores))]
        loss = torch.stack(loss)
        avg_loss = sum(loss)/len(loss)
        losses.append(avg_loss)
losses = torch.stack(losses)
total_loss = sum(losses)/len(losses)  # average over all sentences in batch
total_loss.backward(retain_graph=True)
running_loss += total_loss.item() 
optimizer.step()
count += 1

And code for one possible BaseModel (the others are practically identical):

class LSTMTagger(nn.Module):

def __init__(self, embedding_dim, hidden_dim, dropout, vocab_size, alphabet_size,
             directions=1, device='cpu'):

    super(LSTMTagger, self).__init__()
    self.device = device

    self.hidden_dim = hidden_dim
    self.directions = directions
    self.dropout = nn.Dropout(dropout)

    self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

    # The LSTM takes word embeddings as inputs, and outputs hidden states
    # with dimensionality hidden_dim.
    self.lstm = nn.LSTM(embedding_dim, hidden_dim, dropout=dropout, bidirectional=directions == 2)

    # The linear layer that maps from hidden state space to tag space
    self.hidden = self.init_hidden(hidden_dim)

def init_hidden(self, dim):
    # Before we've done anything, we don't have any hidden state.
    # Refer to the PyTorch documentation to see exactly
    # why they have this dimensionality.
    # The axes semantics are (num_layers, minibatch_size, hidden_dim)
    return (torch.zeros(self.directions, 1, dim).to(device=self.device),
            torch.zeros(self.directions, 1, dim).to(device=self.device))

def forward(self, sentence):
    word_idxs = []
    for word in sentence:
        word_idxs.append(word[0])

    embeds = self.word_embeddings(torch.LongTensor(word_idxs).to(device=self.device))
   
    lstm_out, self.hidden = self.lstm(
        embeds.view(len(sentence), 1, -1), self.hidden)
    lstm_out = self.dropout(lstm_out)
    return lstm_out

Answer 1

The problem is that when I was resetting the hidden states of the model ( model.hidden = model.init_hidden(hidden_dim) ) I didn't actually reassign the reinitialized weights to the BaseModel , but only in the MTLWrapper (which doesn't technically even use hidden layers). I amended my MTLWrapper 's init_hidden() function as follows:

class MTLWrapper(nn.Module):

    def init_hidden(self, hidden_dim):
        self.base_model.hidden = self.base_model.init_hidden(hidden_dim)
        return self.base_model.init_hidden(hidden_dim)

This resolved the first error, and my code runs without the retain_graph=True flag.