with torch.no_grad()

---

合理的利用torch.no_grad()可以节省显存，在推断的时候，可以少占用

模型训练的时候，爆显存了，可以调整batch，对数据进行crop等等操作

今天发现一个模型，训练ok，每次测试的时候爆显存。开始以为是因为用了全图（1920x1080略大）进行inference。

这是一方面，但后来发现忘了用with torch.no_grad():
这导致模型运算的时候不能释放显存（记录了梯度信息），所以显存巨大。
加了之后，用了不过3G显存就够了。。确实inference不需要那么多显存的，以后记着这种不正常现象如何处理。

一般训练不爆显存，测试也不会爆；训练时的显存占用远多于inference

---

 

torch.no_grad() 是一个上下文管理器，被该语句 wrap 起来的部分将不会track 梯度。

例如：

a = torch.tensor([1.1], requires_grad=True)
b = a * 2

• 1
• 2

打印b可看到其 grad_fn 为 mulbackward 表示是做的乘法。

b
Out[63]: tensor([2.2000], grad_fn=<MulBackward0>)

• 1
• 2

b.add_(2)
Out[64]: tensor([4.2000], grad_fn=<AddBackward0>)

• 1
• 2

可以看到不被wrap的情况下，b.grad_fn 为 addbackward 表示这个add 操作被track了

with torch.no_grad():
    b.mul_(2)

• 1
• 2

在被包裹的情况下可以看到 b.grad_fn 还是为 add，mul 操作没有被 track. 但是注意，乘法操作是被执行了的。(4.2 -> 8.4)

b
Out[66]: tensor([8.4000], grad_fn=<AddBackward0>)

• 1
• 2

所以如果有不想被track的计算部分可以通过这么一个上下文管理器包裹起来。这样可以执行计算，但该计算不会在反向传播中被记录。

同时 torch.no_grad() 还可以作为一个装饰器。
比如在网络测试的函数前加上

@torch.no_grad()
def eval():
	...

• 1
• 2
• 3

扩展：
同样还可以用 torch.set_grad_enabled()来实现不计算梯度。
例如：

def eval():
	torch.set_grad_enabled(False)
	...	# your test code
	torch.set_grad_enabled(True)