由于pytorch计算图的构建需要消耗内存和计算资源,在一些情况下,计算图并不是必要的,比如神经网络的推导,在这种情况下,可以使用torch.no_grad 上下文管理器,在这个上下文管理器的作用域里进行的神经网络计算不会构建任何计算图。
另外,还有一种情况是对于一个张量,我们在反向传播的时候可能不需要让梯度通过这个张量的节点,也就是新建的计算图要和原来的计算图分离。在这种情况下,可以使用张量的detach方法,通过调用这个方法,可以返回一个新的张量,该张量会成为一个新的计算图的叶子节点,新的计算图和老的计算图相互分离,互不影响。
具体代码如下:
import torch
torch.cuda.set_device(0)
t1 = torch.randn(3,3,requires_grad=True) # 定义一个3X3的张量
print(t1)
t2 = t1.sum()
print(t2) # t2的计算构建了计算图,输出结果带有grad_fn
with torch.no_grad():
t3 = t1.sum()
print(t3) # t3的计算没有构建计算图,输出结果不带grad_fn
print(t1.sum()) # 保持原来的计算图, 和原来的计算图分离
print(t1.sum().detach())
Done!!!