合理的利用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)