调用pytorch的LogSoftmax后显存暴增

千锤万击 2018-05-29 11:24:57

初次使用pytorch，碰到内存暴增的问题，折腾1个多星期了，依然无法解决（代码附后）。

CIFAR10数据集，代码要完成的任务是：前向传播时会求多个net的输出out，然后求平均。此处（函数forward）每计算一次out，都需要将近7G的显存，多次迭代后显存会成倍增长。而Titan的显存也不过12G，所以在第二次计算out时就提示out of memory，最终导致程序异常终止而无法进行下一步的计算。
针对此问题，我的努力：
1）排除编码错误。因为这是从github上下载的代码，不太可能是编码错误
2）排除版本错误。我以为是pytorch版本导致，结果换了多个版本0.0.12、0.2.0、0.3.0，0.4.0均无法解决此问题
3）排除一次性加载数据导致。因为前面代码中做过
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True, num_workers=4, drop_last=True)
请问这是什么原因导致？有何解决办法？

我的设想：可能是pytorch的动态图导致，每次的计算结果out都是Variable，结果保存在显存中，再计算一次out后，结果依然会保存在显存中。而每次计算的开销又是相当巨大，所以导致显存的成倍增加。是否可以通过下面两种方式解决：
1）将每次前向传播得到的out保存在内存中，在内存中求平均，然后在将平均后的out放入显存，以备后续的反向传播。
2）将out都是的类型由Variable变为Tensor
这两个方法仅仅是设想，不知是否可行，即使可行也不知该如何做。



for i, (images, labels) in enumerate(train_loader):

    images = Variable(images).cuda()

    labels = Variable(labels).cuda()



    out = net(images)

    loss = criterion(out, labels) / UPDATE_EVERY

    loss.backward()



    if (i + 1) % UPDATE_EVERY == 0:

        optimizer.step()

        optimizer.zero_grad()



    loss_epoch = loss_epoch + loss.data[0]

    score_epoch = score_epoch + compute_score(out.data, labels.data)



def forward(self, x):

    out = nn.LogSoftmax()(self.nets[0](x))



    for n in range(1, self.num):

        out = out + nn.LogSoftmax()(self.nets[n](x))



    out = out / self.num

    return out

...全文