Pytorch DataLoader内存未释放

时间:2018-08-25 07:25:27

标签: python machine-learning deep-learning pytorch

我想在Google合作实验室的pythorch上实现SRGAN,但是DataLoader的内存似乎已释放,因此如果您转为纪元,则会发生内存错误。 如果您告诉我如何执行此操作以释放每批内存,将不胜感激。 这是代码的github链接 https://github.com/pacifinapacific/Hello-World/blob/master/Untitled0.ipynb

它变成48,并且在1次回波时发生了内存错误, 如果将批量大小设置为8的1/6,则会在大约6个纪元时出现错误。

我正在使用以下代码读取高分辨率和低分辨率图像。扩展ImageFolder

但是例如,即使执行学习时发生错误,GPU的内存也不会释放

class DownSizePairImageFolder(ImageFolder):
def __init__(self, root, transform=None, large_size=256, small_size=64, **kwds):
    super().__init__(root, transform=transform, **kwds)
    self.large_resizer = transforms.Scale(large_size)
    self.small_resizer = transforms.Scale(small_size)

def __getitem__(self, index):
    path, _ = self.imgs[index]
    img = self.loader(path)
    large_img = self.large_resizer(img)
    small_img = self.small_resizer(img)
    if self.transform is not None:
        large_img = self.transform(large_img)
        small_img = self.transform(small_img)
    return small_img, large_img


train_data = DownSizePairImageFolder('./lfw-deepfunneled/train',   transform=transforms.ToTensor())
test_data = DownSizePairImageFolder('./lfw-deepfunneled/test',    transform=transforms.ToTensor())
batch_size = 8
train_loader = DataLoader(train_data, batch_size, shuffle=True)
test_loader = DataLoader(test_data, batch_size, shuffle=False)

1 个答案:

答案 0 :(得分:1)

每次在模型中传播时,Pytorch都会构建一个计算图。此图通常会保留,直到输出变量G_loss超出范围,例如。当循环中的新迭代开始时。

但是,您将此损失附加到列表中。因此,python仍然知道该变量,并且该图未释放。您可以使用.detach()从当前图形中分离变量(这比我之前建议的.clone()好,因为它也会复制张量的数据)。

作为一个小的副节点:在train()函数中,您在D_loss,G_loss循环中返回for,而不是在循环之后;因此您始终只能使用第一批。