with torch.no_grad() 详解

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

例如：

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

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

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

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

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

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

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

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

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

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

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

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

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