方法一
- 将需要固定的那部分参数的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