时间戳

首页 > 解决方案 > 评估时 GPU 内存不足：Pytorch

pytorch - 评估时 GPU 内存不足：Pytorch

问题描述

当我只训练而不验证时，模型训练得很好，但是它在评估过程中内存不足，但我不明白为什么这可能是一个问题，特别是因为我正在使用torch.no_grad()任何想法？

def test(epoch,net,testloader,optimizer):
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    idx = 0
    features_all = []
    for batch_idx, (inputs, targets) in enumerate(testloader):
        with torch.no_grad():
            idx = batch_idx
            # inputs, targets = inputs.cpu(), targets.cpu()
            if use_cuda:
                inputs, targets = inputs.cuda(), targets.cuda()
            inputs, targets = Variable(inputs), Variable(targets)
            save_features, out, ce_loss = net(inputs,targets)
            test_loss += ce_loss.item()
            _, predicted = torch.max(out.data, 1)
            total += targets.size(0)
            correct += predicted.eq(targets.data).cpu().sum().item()
            features_all.append((save_features, predicted, targets.data))
    test_acc = 100.*correct/total
    test_loss = test_loss/(idx+1)
    logging.info('test, test_acc = %.4f,test_loss = %.4f' % (test_acc,test_loss))
    print('test, test_acc = %.4f,test_loss = %.4f' % (test_acc,test_loss))
    return features_all, test_acc

标签： pytorch

解决方案

features_all.append((save_features, predicted, targets.data))

此行保存对 GPU 内存中张量的引用，因此当循环进入下一次迭代时，CUDA 内存不会被释放（最终导致 GPU 内存不足）。将张量移动到 CPU（使用.cpu()），同时保存它们。

推荐阅读