Why GPU memory usage keeps ceaselessly growing when training the model?

ShawnGuo: Hello everyone. Recently, I implemented a simple recursive neural network. When training this model on sample/small data set, everything works fine. However, when training it on large data and on GPUs, “out of memory” is raised. Along with the training goes on, usage of GPU memory keeps growing up. So, I want to know, why does this happen? I would be grateful if you could help. The model and training procedure are defined as follow: <SCODE>    def train_step(self, data):
        train_loss = 0
        for _data in data:
            p_tree = _data['p_tree']
            h_tree = _data['h_tree']
            if args.cuda:
                target = Variable(torch.LongTensor([_data['label']]).cuda())
            else:
                target = Variable(torch.LongTensor([_data['label']]))
            self.optimizer.zero_grad()
            # self.model is an instance of class RootAlign
            output = self.model(p_tree, h_tree)
            loss = F.nll_loss(output, target)
            loss.backward()
            self.optimizer.step()
            train_loss += loss.data[0]
        return train_loss

class RootAlign(nn.Module):
    def __init__(self, word_embedding, config):
        super(RootAlign, self).__init__()
        self.rnn = VanillaRecursiveNN(word_embedding, config['hidden_dim'], config['cuda_flag'])
        self.linear = nn.Linear(config['hidden_dim'] * 2, config['relation_num'])

    def forward(self, p_tree, h_tree):
        p_tree.postorder_traverse(self.rnn)
        h_tree.postorder_traverse(self.rnn)

        out = F.log_softmax(self.linear(F.sigmoid(torch.cat((p_tree.calculate_result, h_tree.calculate_result), 1))))
        return out

class VanillaRecursiveNN(nn.Module):
    def __init__(self, word_embedding, hidden_dim, cuda_flag=False):
        super(VanillaRecursiveNN, self).__init__()
        self.word_dim = word_embedding.embeddings.size(1)
        self.hidden_dim = hidden_dim

        self.embedding = nn.Embedding(word_embedding.embeddings.size(0),
                                      self.word_dim)
        self.embedding.weight = nn.Parameter(word_embedding.embeddings)

        self.word2hidden = nn.Linear(self.word_dim, self.hidden_dim, False)
        self.hidden2hidden = nn.Linear(2 * self.hidden_dim, self.hidden_dim)

        self.cuda_flag = cuda_flag

    def forward(self, node):
        if not node.val is None:
            if self.cuda_flag:
                node.calculate_result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]).cuda())))
            else:
                node.calculate_result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]))))
            return node.calculate_result
        else:
            assert len(node.children) == 2
            node.calculate_result = self.hidden2hidden(torch.cat((node.children[0].calculate_result,
                                                          node.children[1].calculate_result), 1))
            return node.calculate_result
<ECODE>

apaszke: 

OCY: I observed similar GPU memory behavior. <SCODE>def function1(A):
    B = A**2 - 2*A
    C = torch.sqrt(B)
    return C

def function2(A):
    return torch.sqrt(A**2 - 2*A)
<ECODE> Thus, in terms of memory usage, it seems that function2 is twice efficient than function1. This doesn’t apply to all the cases, but in many cases, removing intermediate variables reduces GPU memory usage a lot in my programs. This seems to be that underlying CUDA does not free memory immediately after the moment that memory is not needed anymore. I think some GPU memory Garbage Collection method in pytorch is needed for efficient GPU memory management.

apaszke: Both functions will consume the same amount of memory. The execution will look like this (in parenthesis you have current/peak memory usage in multiplies of A size): Assume A is allocated (1/1) Compute A**2 (2/2) Compute 2*A (3/3) Compute A**2 - 2*A (4/4) Free A**2 and 2*A (2/4) Compute torch.sqrt(B) (3/4) Return and free everything except the input and result (2/4) There’s no way for the framework to know when a tensor won’t be needed anymore, we don’t have that knowledge upfront, abnd this is why it’s impossible to implement any garbage collection. The memory management is already very efficient and all tensors are freed as soon as you let them go.

OCY: Thanks! Things became clear.

ShawnGuo: Thanks for the help. I did the replacement as you stated above in the following ways. <SCODE># way No.1
if not node.val is None:
            if self.cuda_flag:
                variable =Variable(torch.LongTensor([node.word_id]).cuda())
            else:
                variable = Variable(torch.LongTensor([node.word_id]))
            result = self.word2hidden(self.embedding(variable))
            node.calculate_result = result.data
            return node.calculate_result

# way No.2
if not node.val is None:
            if self.cuda_flag:
                node.calculate_result = self.word2hidden(self.embedding(
                        Variable(torch.LongTensor([node.word_id]).cuda()))).data
            else:
                node.calculate_result = self.word2hidden(self.embedding( 
                        Variable(torch.LongTensor([node.word_id])))).data
            return node.calculate_result
# way No.3
if not node.val is None:
            if self.cuda_flag:
                result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]).cuda())))
            else:
                result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]))))
            node.calculate_result = result.data
            return node.calculate_result
<ECODE> <SCODE>Traceback (most recent call last):
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 172, in <module>
    t.train()
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 111, in train
    train_loss = self.train_step(self.data.train)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 143, in train_step
    output = self.model(p_tree, h_tree)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 210, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/align_model.py", line 18, in forward
    p_tree.postorder_traverse(self.rnn)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 187, in postorder_traverse
    func(self)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 219, in __call__
    var = var[0]
TypeError: 'float' object has no attribute '__getitem__'
<ECODE> <SCODE>Traceback (most recent call last):
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 172, in <module>
    t.train()
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 111, in train
    train_loss = self.train_step(self.data.train)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/trainer.py", line 143, in train_step
    output = self.model(p_tree, h_tree)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 210, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/align_model.py", line 18, in forward
    p_tree.postorder_traverse(self.rnn)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 186, in postorder_traverse
    c.postorder_traverse(func)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree.py", line 187, in postorder_traverse
    func(self)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 210, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/shawnguo/PythonWS/KnowledgeEnhancedTE/tree_models.py", line 26, in forward
    self.embedding(Variable(torch.LongTensor([[node.word_id]]).cuda())))
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 210, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/linear.py", line 52, in forward
    return self._backend.Linear()(input, self.weight)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/_functions/linear.py", line 10, in forward
    output.addmm_(0, 1, input, weight.t())
RuntimeError: matrix and matrix expected at /home/soumith/local/builder/wheel/pytorch-src/torch/lib/THC/generic/THCTensorMathBlas.cu:235
<ECODE> Anyway, thanks again for your help. Looking forward for your reply.

jhjungCode: I think it’s just wrong input number of fc-layer and usage of cat() fucntion <SCODE>import torch
hidden_dim = 10

x = torch.randn(hidden_dim, 1).cuda()
print(x.size())
y = torch.cat((x, x), 1)
print(y.size())
y = torch.cat((x, x), 0)
print(y.size())
<ECODE> result is below so, fixed code is <SCODE>self.hidden2hidden = nn.Linear(2 * self.hidden_dim, self.hidden_dim)
...
node.calculate_result = self.hidden2hidden(torch.cat((node.children[0].calculate_result, node.children[1].calculate_result), 0))
<ECODE>

ShawnGuo: Thanks for your help. However, the problem occurs in the following code: <SCODE>if not node.val is None:
            if self.cuda_flag:
                node.calculate_result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]).cuda())))
            else:
                node.calculate_result = self.word2hidden(
                    self.embedding(Variable(torch.LongTensor([node.word_id]))))
            return node.calculate_result
<ECODE> And, I’ve found a puzzling phenomenon, if the above code have been changed to: <SCODE>if not node.val is None:
            if self.cuda_flag:
                variable =Variable(torch.LongTensor([node.word_id]).cuda())
            else:
                variable = Variable(torch.LongTensor([node.word_id]))
            node.calculate_result = self.word2hidden(self.embedding(variable))
            return node.calculate_result 
<ECODE>

jhjungCode: You should check the connectivity between network layers Try it <SCODE>self.embedding(Variable(torch.LongTensor(node.word_id).cuda())))
<ECODE> or add squeeze variable <SCODE>variable = variable.squeeze()
<ECODE>

ShawnGuo: Anyway, thanks for your help.

apaszke: You don’t need to save anything for backprop, autograd will take care of that, and my solution is valid. The problems you’re having are only due to giving inputs of invalid sizes to different modules. You can print them inside your module and see if they are what you expect, and what matches the requirements specified in the docs.

ShawnGuo: Yes，you’re right. The computation in case “not node.val is None” is correct. Problem is in the computation of the other case. I’m trying to fix it. Thank you very much!

ShawnGuo: <SCODE>if not node.val is None:
            if self.cuda_flag:
                variable =Variable(torch.LongTensor([node.word_id]).cuda())
            else:
                variable = Variable(torch.LongTensor([node.word_id]))
            result = self.word2hidden(self.embedding(variable))
            node.calculate_result = result.data
            return node.calculate_result
<ECODE> <SCODE>while not isinstance(var, Variable):
        var = var[0]
<ECODE> After I do the following change, everything works again except the weights of model hasn’t been updated. No.1, change ni RootAlign: <SCODE>class RootAlign(nn.Module):
    def __init__(self, word_embedding, config):
        super(RootAlign, self).__init__()
        self.rnn = VanillaRecursiveNN(word_embedding, config['hidden_dim'], config['cuda_flag'])
        self.linear = nn.Linear(config['hidden_dim'] * 2, config['relation_num'])

    def forward(self, p_tree, h_tree):
        p_tree.postorder_traverse(self.rnn)
        h_tree.postorder_traverse(self.rnn)

        p_result = Variable(p_tree.calculate_result)
        h_result = Variable(h_tree.calculate_result)
        out = F.log_softmax(self.linear(F.sigmoid(
            torch.cat((p_result, h_result), 1))))
        return out
<ECODE> No.2 Change in VanillarRecursiveNN: <SCODE>def forward(self, node):
    if not node.val is None:
        if self.cuda_flag:
            result = self.word2hidden(self.embedding(
                Variable(torch.LongTensor([node.word_id]).cuda())))
        else:
            result = self.word2hidden(self.embedding(
                Variable(torch.LongTensor([node.word_id]))))
        node.calculate_result = result.data
        return result
    else:
        assert len(node.children) == 2
        l_result = Variable(node.children[0].calculate_result)
        r_result = Variable(node.children[1].calculate_result)
        result = self.hidden2hidden(torch.cat((l_result, r_result), 1))
        node.calculate_result = result.data
        return result
<ECODE>

apaszke: 

ShawnGuo: Well, I guess that this is the key to address the problem. I’ll try immediately.

ShawnGuo: I have an additional question. How to batch tree data when training model? Every tree has their own structure. How can I batch them under the current implementation of Recursive model? ps: The time of every epoch is about 4.5 hr on SNLI.(GPU: Titan X) It takes too long.

jekbradbury: 

ShawnGuo: Your implementation is cool~ I’ll learn from it and try to batch data on my model. Thanks very much.