GPU 出 memory 评估：Pytorch

Question

The model trains fine when I only train and don't validate, however it runs out of memory during evaluation, but I don't understand why this might be a problem especially since I am using torch.no_grad() any ideas?

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

Answer 1

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

This line is saving references to tensors in GPU memory and so the CUDA memory won't be released when loop goes to next iteration (which eventually leads to the GPU running out of memory). Move the tensors to CPU (using .cpu() ) while saving them.