[Pytorch, timm] Optimizer & Parameter Group Learning Rate

pytorch 에서 일반적으로 optimizer 를 다음과 같이 사용

optimizer = optim.Adam(model, lr=0.0001)

 

만약 model 의 각 parameter 에 다른 옵션 (learning rate, eps 등)을 주고 싶은 경우 다음과 같이 각 parameter group 지정

optim.Adam([{'params': model.base.parameters()},
            {'params': model.classifier.parameters(), 'lr': 1e-3}], lr=1e-2)

 

Parameter group 은 다음으로 확인 가능

#n 번째 parameter group 을 보고 싶은 경우
print(optimizer.param_groups[n]) 

>>> Adam (Parameter Group n
          amsgrad: False
          betas: (0.9, 0.999)
          capturable: False
          eps: 1e-08
          foreach: None
          lr: 0.0002
          maximize: False
          weight_decay: 0.0)

 

만약 특정 paramter group 에서 일부 옵션을 수정하고 싶은 경우

# Parameter group 은 dictionary 로 구성되어 있음
optimizer.param_groups[n]['lr'] = 0.03

print(optimizer.param_groups[n])

>>> Adam (Parameter Group n
          amsgrad: False
          betas: (0.9, 0.999)
          capturable: False
          eps: 1e-08
          foreach: None
          lr: 0.03 # 수정됨
          maximize: False
          weight_decay: 0.0)

 

하지만 timm(pytorch image models) 라이브러리에선 일반적인 pytorch optimizer 생성과는 다름

def create_optimizer_v2(
        model_or_params,
        opt: str = 'sgd',
        lr: Optional[float] = None,
        weight_decay: float = 0.,
        momentum: float = 0.9,
        filter_bias_and_bn: bool = True,
        layer_decay: Optional[float] = None,
        param_group_fn: Optional[Callable] = None,
        **kwargs):
    """ Create an optimizer.

    TODO currently the model is passed in and all parameters are selected for optimization.

맨 밑 줄을 보면 알 수 있듯, 모든 parameter 가 한 개의 parameter group으로 묶임

 

모델이 encoder 와 classifier 로 구성이 되어 있어도 optimizer 를 print 해보면 Parameter Group 0 으로 하나만 출력됨

 

따라서 일단 create_optimizer 를 통해 특정 parameter 에 대한 optimizer 를 생성한 후, 다른 parameter group 의 dictionary 를 추가해줘야함 → Add param group 사용

 

optimizer = create_optimizer(args, encoder)
cls_dict = copy.deepcopy(optimizer.param_groups[0])
cls_dict['params'] = model.classifier.parameters()
cls_dict['lr'] = args.cls_lr
optimizer.add_param_group(cls_dict)