DeeplabV3+ 모델 전이 학습(transfer learning) 쉽게 구현하기

 

아래 repository를 기반으로 작성하였다.

GitHub - qubvel/segmentation_models.pytorch: Segmentation models with pretrained backbones. PyTorch.

 

1. Pretrained 모델 로드

# train.py

pretrained_model = smp.DeepLabV3Plus(
    encoder_name='resnet50', 
    encoder_weights=None, 
    in_channels=3, 
    classes=1
)
pretrained_model.load_state_dict(torch.load('pretrained_model.pth'))
pretrained_model.to(device)

encoder_weights=None으로 해야 한다. (default 값은 'imagenet')

 

2. (Optional) Freezing Layers

# Freeze all layers
for param in pretrained_model.parameters():
    param.requires_grad = False

# Unfreeze the last few layers (e.g., the classifier layers)
for param in pretrained_model.decoder.parameters():
    param.requires_grad = True

모델의 앞 쪽(lower) layer에서 보다 더 일반적인 특징들(edge나 texture 같은)을 더 잘 학습한다.
따라서 이 부분의 weights는 고정시켜두고 나머지 뒷 쪽(higher) layer만 업데이트하는 것을 고려해볼 수 있다.

 

특히 새로 학습하고자 하는 데이터셋의 크기가 작을 경우에는 이렇게 classifier 부분만 업데이트하는 것이 유리할 수 있다.

https://velog.io/@xuio/Transfer-Learning%EA%B3%BC-Fine-Tuning

위 그림처럼 데이터셋 크기와 데이터셋 또는 task의 유사도에 따라 오른쪽과 같은 fine-tuning 방식을 고민해볼 수 있다.

만약 lower layer도 일부 학습하도록 하고 싶다면 아래와 같이 원하는 layer 개수를 입력하여 구현하면 된다.

# Define the number of layers you want to unfreeze in the encoder. 
# The last 'n' layers of the encoder can be unfrozen for training.
n = 2 # You can define your value of n here

# Unfreeze the last 'n' layers in the encoder
for param in list(pretrained_model.encoder.parameters())[-n:]:
    param.requires_grad = True

얼마나 layer를 고정하는 것이 좋은지는 약간의 실험이 필요하다.
우선 몇 개 layer를 고정하는 것부터 시작해보고, 필요할 경우 고정하려는 layer 개수를 조정하는 것이 좋다.

또한 fine-tuning할 때에는 모든 layer를 학습시킬 때보다 learning rate이 조금 더 작은 것이 좋다.

코드를 전체적으로 살펴보면 다음과 같다.

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = DeepLabV3Plus(
        encoder_name='resnet50',
        encoder_weights=None,
        in_channels=3,
        classes=1
    )
    model.load_state_dict(torch.load('pretrained_crack_model.pth'))
    
    for param in model.parameters():  # Freeze all layers
        param.requires_grad = False
    for param in model.decoder.parameters():  # Unfreeze the classifier layers
        param.requires_grad = True

    model.to(device)

    epochs = 1000
    loss_fn = DiceLoss('binary', from_logits=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-3)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, verbose=True)
    
    # dataset load -> run train()