[Pytorch] 학습된 모델에서 필요한 layer 만 사용하는 방법

import torch.nn as nn
import torchvision.models as models

# Pre-trained ResNet 모델 불러오기
model = models.resnet50(pretrained=True)

# ResNet50에서 마지막 classification layer를 제거합니다.
modules = list(model.children())[:-1]
model = nn.Sequential(*modules)

# fully connected layer 1층 추가
model.fc = nn.Linear(2048, 1024)

# 학습된 모델을 evaluation mode로 설정합니다.
model.eval()

위 코드에서는 먼저 PyTorch에서 제공하는 torchvision.models.resnet50 함수를 사용하여 ResNet50 모델을 불러옵니다. 다음으로, 마지막 classification layer를 제거하기 위해 ResNet50 모델의 children을 모두 가져온 후 마지막 layer를 제외한 모든 layer들을 리스트에서 제거합니다. 그 다음, nn.Sequential() 함수를 사용하여 새로운 모델을 만듭니다. 이 새로운 모델은 ResNet50 모델에서 classification layer를 제거한 feature extractor 모델입니다.

그 다음으로, fully connected layer 1층을 추가하기 위해 model.fc를 nn.Linear(2048, 1024)로 변경합니다. 여기서 2048은 ResNet50 모델에서 마지막 feature map의 크기이며, 1024는 출력 벡터의 크기입니다.

마지막으로, 학습된 모델을 evaluation mode로 설정합니다. 이렇게 하면 학습된 가중치를 고정하여 모델을 사용할 수 있습니다.