实际上,detach()
就是返回一个新的tensor,并且这个tensor是从当前的计算图中分离出来的。但是返回的tensor和原来的tensor是共享内存空间的。
import torch
a = torch.tensor([1.0, 2.0, 3.0], requires_grad = True)
a = a.detach() # 会将requires_grad 属性设置为False
print(a.requires_grad)
结果为:
False
举个例子来说明一下detach()
函数有什么用。 如果A网络的输出被喂给B网络作为输入, 如果我们希望在梯度反传的时候只更新B中参数的值,而不更新A中的参数值,这时候就可以使用detach()
。
a = A(input)
a = a.deatch() # 或者a.detach_()进行in_place操作
out = B(a)
loss = criterion(out, labels)
loss.backward()
如果希望修改A的参数, 而不希望修改B的参数, 那么就需要手动将B中参数的requires_grad属性设置为False:
for param in B.parameters():
param.requires_grad = False