# 1. 只训练部分层
class RESNET_attention(nn.Module): | |
def __init__(self, model, pretrained): | |
super(RESNET_attetnion, self).__init__() | |
self.resnet = model(pretrained) | |
for p in self.parameters(): | |
p.requires_grad = False | |
self.f = nn.Conv2d(2048, 512, 1) | |
self.g = nn.Conv2d(2048, 512, 1) | |
self.h = nn.Conv2d(2048, 2048, 1) | |
self.softmax = nn.Softmax(-1) | |
self.gamma = nn.Parameter(torch.FloatTensor([0.0])) | |
self.avgpool = nn.AvgPool2d(7, stride=1) | |
self.resnet.fc = nn.Linear(2048, 10) |
- 这样就将 for 循环以上的参数固定，只训练下面的参数。但是注意需要在 optimizer 中添加上这样的一句话 filter (lambda p: p.requires_grad, model.parameters ()。
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5) |
# 2. 固定部分层参数
for k,v in model.named_parameters(): | |
if k!='XXX': | |
v.requires_grad=False#固定参数 |
# 3. 检查部分参数是否固定
for k,v in model.named_parameters(): | |
if k!='xxx.weight' and k!='xxx.bias' : | |
print(v.requires_grad)#理想状态下，所有值都是 False |
# 4. 查看可训练参数
for name, param in model.named_parameters(): | |
if param.requires_grad: | |
print(name) |
# 5. 查看网络总参数
net = Model() | |
print('# Model parameters:', sum(param.numel() for param in net.parameters())) |
