PyTorch冻结网络参数

方法一

  • 将需要固定的那部分参数的requires_grad置为False.
  • 在优化器中加入filter根据requires_grad进行过滤.
# requires_grad置为False
for p in net.XXX.parameters():
    p.requires_grad = False

# filter
optimizer.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)

方法二

这种方式只需要在网络定义中的forward方法中,将需要冻结的层放在 torch.no_grad()下,强力推这种方式。

class xxnet(nn.Module):
    def __init__():
        ....
        self.layer1 = xx
        self.layer2 = xx
        self.fc = xx

    def forward(self.x):
        with torch.no_grad():
            x = self.layer1(x)
            x = self.layer2(x)
        x = self.fc(x)
        return x

猜你喜欢

转载自blog.csdn.net/weixin_42990464/article/details/115265755