通过PyTorch训练时记忆不断积累

Question

我正在使用PyTorch训练深度学习模型。由于未知原因，内存不断累积，导致会话在30个时期内被杀死并不足。

这里有些想法：

1. 想知道是否是由matplotlib引起的，所以我添加了plt.close('all')；没用

2. 添加了gc.collect()；没用

3. 想知道这是否是由cv2.imwrite()引起的，但不知道如何检查。有什么建议吗？

4. PyTorch问题？

5. 其他...

   model.train()
   for epo in range(epoch):
       for i, data in enumerate(trainloader, 0):
           inputs = data
           inputs = Variable(inputs)
           optimizer.zero_grad()
   
           top = model.upward(inputs + white(inputs))
           outputs = model.downward(top, shortcut = True)
   
   
           loss = criterion(inputs, outputs)
           loss.backward()
           optimizer.step()
   
           # Print generated pictures every 100 iters
           if i % 100 == 0:
               inn = inputs[0].view(128, 128).detach().numpy() * 255
               cv2.imwrite("/home/tk/Documents/recover/" + str(epo) + "_" + str(i) + ".png", inn)
   
               out = outputs[0].view(128, 128).detach().numpy() * 255
               cv2.imwrite("/home/tk/Documents/recover/" + str(epo) + "_" + str(i) + "_re.png", out)
   
           # Print loss every 50 iters
           if i % 50 == 0:
               print ('[%d, %5d] loss: %.3f' % (epo, i, loss.item()))
   
       gc.collect()
       plt.close("all")
   

================================================ ====================

20181222更新

数据集和DalaLoader

class MSourceDataSet(Dataset):

    def __init__(self, clean_dir):



        for i in cleanfolder:
            with open(clean_dir + '{}'.format(i)) as f:
                clean_list.append(torch.Tensor(json.load(f)))


        cleanblock = torch.cat(clean_list, 0)
        self.spec = cleanblock


    def __len__(self):
        return self.spec.shape[0]


    def __getitem__(self, index): 

        spec = self.spec[index]
        return spec

trainset = MSourceDataSet(clean_dir)
trainloader = torch.utils.data.DataLoader(dataset = trainset,
                                            batch_size = 4,
                                            shuffle = True)

模型真的很复杂而且很长...再加上以前（使用相同模型）没有发生内存累积问题，所以我不会在这里发布它...