with torch.no_grad():显著减少测试时显存占用

问题描述

        将训练好的模型拿来做inference，发现显存被占满，无法进行后续操作，但按理说不应该出现这种情况。

RuntimeError: CUDA out of memory. Tried to allocate 128.00 MiB (GPU 0; 7.93 GiB total capacity; 6.94 GiB already allocated; 10.56 MiB free; 7.28 GiB reserved in total by PyTorch)

解决方案

        经过排查代码，发现做inference时，各模型虽然已经设置为eval()模式，但是并没有取消网络生成计算图这一操作，这就导致网络在单纯做前向传播时也生成了计算图，从而消耗了大量显存。

        所以，将模型前向传播的代码放到with torch.no_grad()下，就能使pytorch不生成计算图，从而节省不少显存

with torch.no_grad()：
    # 代码块
    outputs = model(inputs)
	# 代码块

        经过修改，再进行inference就没有遇到显存不够的情况了。此时显存占用显著降低，只占用5600MB左右（3卡）。

model.eval()和torch.no_grad()

model.eval()

• 使用model.eval()切换到测试模式，不会更新模型的k，b参数
• 通知dropout层和batchnorm层在train和val中间进行切换在。train模式，dropout层会按照设定的参数p设置保留激活单元的概率（保留概率=p，比如keep_prob=0.8），batchnorm层会继续计算数据的mean和var并进行更新。在val模式下，dropout层会让所有的激活单元都通过，而batchnorm层会停止计算和更新mean和var，直接使用在训练阶段已经学出的mean和var值
• model.eval()不会影响各层的gradient计算行为，即gradient计算和存储与training模式一样，只是不进行反向传播(backprobagation)，即只设置了model.eval()pytorch依旧会生成计算图，占用显存，只是不使用计算图来进行反向传播。

torch.no_grad()

首先从requires_grad讲起：

requires_grad
        在pytorch中，tensor有一个requires_grad参数，如果设置为True，则反向传播时，该tensor就会自动求导，并且保存在计算图中。tensor的requires_grad的属性默认为False,若一个节点（叶子变量：自己创建的tensor）requires_grad被设置为True，那么所有依赖它的节点requires_grad都为True（即使其他相依赖的tensor的requires_grad = False）

        当requires_grad设置为False时,反向传播时就不会自动求导了，也就不会生成计算图，而GPU也不用再保存计算图，因此大大节约了显存或者说内存。

with torch.no_grad

        在该模块下，所有计算得出的tensor的requires_grad都自动设置为False。

        即使一个tensor（命名为x）的requires_grad = True，在with torch.no_grad计算，由x得到的新tensor（命名为w-标量）requires_grad也为False，且grad_fn也为None,即不会对w求导。例子如下所示：

x = torch.randn(10, 5, requires_grad = True)
y = torch.randn(10, 5, requires_grad = True)
z = torch.randn(10, 5, requires_grad = True)
with torch.no_grad():
    w = x + y + z
    print(w.requires_grad)
    print(w.grad_fn)
print(w.requires_grad)


False
None
False

        也就是说，在with torch.no_grad结构中的所有tensor的requires_grad属性会被强行设置为false，如果前向传播过程在该结构中，那么inference过程中都不会产生计算图，从而节省不少显存。

版本问题

问题描述

volatile was removed and now has no effect. Use with torch.no_grad(): instead

源代码

captions = Variable(torch.from_numpy(captions), volatile=True)

原因

1. 在torch版本中volatile已经被移除。在pytorch 0.4.0之前 input= Variable(input, volatile=True) 设置volatile为True ，只要是一个输入为volatile，则输出也是volatile的，它能够保证不存在中间状态；但是在pytorch 0.4.0之后取消了volatile的机制，被替换成torch.no_grad()函数
2. torch.no_grad() 是一个上下文管理器。在使用pytorch时，并不是所有的操作都需要进行计算图的生成（计算过程的构建，以便梯度反向传播等操作）。而对于tensor的计算操作，默认是要进行计算图的构建的，在这种情况下，可以使用 with torch.no_grad():，强制之后的内容不进行计算图构建。在torch.no_grad() 会影响pytorch的反向传播机制，在测试时因为确定不会使用到反向传播因此 这种模式可以帮助节省内存空间。同理对于 torch.set_grad_enable(grad_mode)也是这样

参考链接

【pytorch系列】 with torch.no_grad():用法详解_sazass的博客-优快云博客

 深入理解model.eval()与torch.no_grad()_雷恩Layne-优快云博客_torch.no_grad()是什么

解决方案：volatile was removed and now has no effect. Use `with torch.no_grad():` instead._修炼清爽的博客-优快云博客