控制自动求导autograd

Contact me

本系列博客主页及相关见此处

来自stackoverflow。

requires_grad=False

如果想冻结部分模型，只训练其他部分, 可以设置这部分的requires_grad = False。

model = torchvision.models.vgg16(pretrained=True)
for param in model.features.parameters():
    param.requires_grad = False

这时，中间的导数缓存就不保存了。

torch.no_grad()

这种情况下，计算的结果都会设置为requires_grad=False, 即便是输入为requires_grad=True。但是就不能反向传播到更前面有梯度的部分了，例如:

x = torch.randn(2, 2)
x.requires_grad = True

lin0 = nn.Linear(2, 2)
lin1 = nn.Linear(2, 2)
lin2 = nn.Linear(2, 2)
x1 = lin0(x)
with torch.no_grad():    
    x2 = lin1(x1)
x3 = lin2(x2)
x3.sum().backward()
print(lin0.weight.grad, lin1.weight.grad, lin2.weight.grad)

输出为：

(None, None, tensor([[-1.4481, -1.1789],
         [-1.4481, -1.1789]]))

lin0.weight.requires_grad是True, 但是没有梯度。

model.eval()

如果不是为了微调，而是仅用于预测，最有效的方法是使用torch.no_grad。这时候也需要设置eval():

model = torchvision.models.vgg16(pretrained=True)
model.eval()

它会设置模型内部self.training为False, 但是也会改变像Dropout或者是BatchNorm的行为，导致训练和测试阶段结果不一致。