Pytorch学习笔记：根据optimizer_type选择优化器

---

1.根据optimizer_type选择优化器

代码如下：

#---------------------------------------#
#   根据optimizer_type选择优化器
#---------------------------------------#
pg0, pg1, pg2 = [], [], []  
for k, v in model.named_modules():
    if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
        pg2.append(v.bias)    
    if isinstance(v, nn.BatchNorm2d) or "bn" in k:
        pg0.append(v.weight)    
    elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
        pg1.append(v.weight)   
optimizer = {
    'adam'  : optim.Adam(pg0, Init_lr_fit, betas = (momentum, 0.999)),
    'sgd'   : optim.SGD(pg0, Init_lr_fit, momentum = momentum, nesterov=True)
}[optimizer_type]
optimizer.add_param_group({"params": pg1, "weight_decay": weight_decay})
optimizer.add_param_group({"params": pg2})

2.pytorch中冻结某些层设置

代码如下：

from torchvision import models
model = models.resnet18()
# 冻结layer4以前的所有层，只训练layer4与全连接层（根据类别需自定义）
no_freeze = ['layer4','fc']  # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
    v.requires_grad = True  # train all layers
    if not any(x in k for x in no_freeze):
        print('freezing %s' % k)
        v.requires_grad = False

---