[pytorch] pretrained model 쉽게 사용하는 방법

이번 포스팅에서는 사전학습 모델(pretrained model)을 사용하는 법을 알아보자.

사전학습 모델(pretrained model)은 torchvision에서 기본 제공한다.

 

FAQ.
1. 전이 학습이란?
다른 데이터셋으로 잘학습되어 있는 모델을 사용
장점: 이미 다량의 데이터로 학습되어 있기 때문에 데이터 부족 문제 해결

2. 학습된 데이터셋이 우리가 사용하려는 데이터셋과 전혀 다르다면 효과가 없지 않나?
데이터 셋과 연관도가 낮다면 상대적으로 성능은 떨어질 수 있다. 하지만 가지고 있는 데이터가 적다면 pretrained 모델을 사용하는 것이 더 좋다.

 

본 포스팅은 아래 순서로 진행된다.

 

1. 사용 가능 모델 확인 

2. 모델 로드하기

3. 모델변경해보기(fine-tuning)

4. transform 확인

 

필요 모듈 설치

pip install torchvision
pip install torchinfo

 

사용 가능 모델 확인

  torchvision  에서 제공하는 기본 모델이 상당히 많다.

import torchvision

for name in dir(torchvision.models):
    print(name)

# output
'''
AlexNet
AlexNet_Weights
ConvNeXt
ConvNeXt_Base_Weights
ConvNeXt_Large_Weights
ConvNeXt_Small_Weights
ConvNeXt_Tiny_Weights
DenseNet
DenseNet121_Weights
DenseNet161_Weights
DenseNet169_Weights
DenseNet201_Weights
EfficientNet
EfficientNet_B0_Weights
EfficientNet_B1_Weights
EfficientNet_B2_Weights
EfficientNet_B3_Weights
EfficientNet_B4_Weights
EfficientNet_B5_Weights
EfficientNet_B6_Weights
EfficientNet_B7_Weights
EfficientNet_V2_L_Weights

...
'''

  dir()   함수를 이용하면 사용 가능 모델을 확인할 수 있다.

 

_Weights가 있는 것도 있고 없는 것도 있다.

_Weights는 해당 모델 파라미터값이라고 보면 되고 _Weights가 없는 것은 모델이다.

 

pretrained 모델 로드

1) 옛날 방법

 

먼저 조건을 걸어 _(suffix)없는 녀석들만 추린다.

즉 weights파일이 아닌 모델 structure만 뽑아내겠다는 거다.

 

for name in dir(torchvision.models):
  if ("Weight" not in name) and ("__" not in name) and (not name.startswith('_')):
    print(name)

# output
'''
AlexNet
ConvNeXt
DenseNet
EfficientNet
GoogLeNet
GoogLeNetOutputs
Inception3
InceptionOutputs
MNASNet
MaxVit
MobileNetV2
MobileNetV3
RegNet
ResNet
ShuffleNetV2
SqueezeNet
SwinTransformer
VGG
VisionTransformer
alexnet
convnext
convnext_base
convnext_large
convnext_small
convnext_tiny
densenet
densenet121
densenet161
densenet169
densenet201
detection
efficientnet
efficientnet_b0
efficientnet_b1
efficientnet_b2
efficientnet_b3
...........
'''

 

예제에선   efficientent_b0  을 기준으로 사용하겠다.

참고.
efficientnet_b0 => efficientnet_b1 => efficientnet_b2 => ...
0부터 숫자가 커질수록 성능은 좋으나 모델이 무거워진다로 이해하면 된다.

 

  pretrained=True   는 pretrained 된 weight를 사용하겠다는 의미이다.

model_loaded_deprecated = torchvision.models.efficientnet_b0(pretrained=True)

 

2) 최신방법(torchvision v0.13+ 이후)

weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
model_loaded = torchvision.models.efficientnet_b0(weights=weights)

weights 파일을 할당받은 후 그 값을 efficientnet_b0 모델 인자로 넣는 형태다.

 

.DEFAULT값은 EfficientNet_B0_Weights.IMAGENET1K_V1를 가리킨다.

print(torchvision.models.EfficientNet_B0_Weights.DEFAULT)
# EfficientNet_B0_Weights.IMAGENET1K_V1

 

최신 방법은 짤 짜인 모델 structure를 그대로 이용하면서 customized weight를 사용할 수도 있다.

torchvision.models.efficientnet_b0(weights=torch.load('모델weights.pth'))

물론  customized weight는 모델과 호환이 되어야 성능이 나온다.

 

모델 변경

1) layer freeze하기

학습이 안되도록 꽁꽁 얼린다는 표현으로 freeze라고 한다.

 

아래 방법은 모델에 있는 모든 파라미터를 꽁꽁 얼려 학습이 안되도록 하는 코드다.

for param in model_loaded.features.parameters(): 
    param.requires_grad = False

inference용으로 사용한다면 위와 같이 작성해도 문제가 없다.

 

반면 전체 layer가 아닌 특정 layer만 얼리고 싶다면 어떻게 할까?

  model_loaded.named_parameters()  를 사용하면 각 layer의 이름과 맵핑된 파라미터값을 알 수 있다.

for name, param in model_loaded.named_parameters():
    print('name :',name)

# output
'''
name : features.0.0.weight
name : features.0.1.weight
name : features.0.1.bias
name : features.1.0.block.0.0.weight
name : features.1.0.block.0.1.weight
name : features.1.0.block.0.1.bias
name : features.1.0.block.1.fc1.weight
name : features.1.0.block.1.fc1.bias
name : features.1.0.block.1.fc2.weight
name : features.1.0.block.1.fc2.bias
name : features.1.0.block.2.0.weight
name : features.1.0.block.2.1.weight
name : features.1.0.block.2.1.bias
name : features.2.0.block.0.0.weight
name : features.2.0.block.0.1.weight
name : features.2.0.block.0.1.bias
name : features.2.0.block.1.0.weight
name : features.2.0.block.1.1.weight
name : features.2.0.block.1.1.bias
name : features.2.0.block.2.fc1.weight
name : features.2.0.block.2.fc1.bias
name : features.2.0.block.2.fc2.weight
name : features.2.0.block.2.fc2.bias
name : features.2.0.block.3.0.weight
name : features.2.0.block.3.1.weight
name : features.2.0.block.3.1.bias
name : features.2.1.block.0.0.weight
							.
							.
							.
							.
							.

 

원하는 parameter name을 찾아 조건문을 걸면 특정 파라미터 freeze여부를 결정할 수 있다.

for name, param in model_loaded.named_parameters():
    if name == 'features.2.0.block.3.0.weight':
        param.requires_grad = False

 

역으로 freeze를 거는 게 아니라 학습시킬 파라미터만 추려내는 방법도 있다.

학습시킬 param을 리스트에 담고 그 값을 optimizer 인자로 넣어준다.

parameters_to_update = []
for name, param in model_loaded.named_parameters():
    if 'layer1' in name or 'layer2' in name:
        parameters_to_update.append(param)

optimizer = optim.SGD(parameters_to_update, lr=0.001, momentum=0.9)

일반적으로 optimizer 첫 번째 인자에는   model.parameters()  가 들어간다.
이는 전체 파라미터를 학습하는데 쓰겠다는 의미다.

 

2) 모델 구조(structure) 바꾸기

지금까지 모델의 학습 가능 여부를 조정하였다면 모델 구조 자체를 변경하는 방법에 대해 알아보자.

 

CNN 모델을 fine-tuning 할 때 마지막   output_feature   개수를 변경하는 경우가 많다.

그래서 그런지 파이토치는   .classifier  라는 프러퍼티를 특별하게 제공해 준다.

(모델에 따라   .classifier  가아닌   .fc  로 접근해야 할수도 있으니 참고하기 바란다.)

print(model_loaded.classifier)
# output

'''
Sequential(
  (0): Dropout(p=0.2, inplace=True)
  (1): Linear(in_features=1280, out_features=1000, bias=True)
)
'''

  EfficientNet   모델은 imagenet(1000개 클래스)으로 pretrained 된 모델이다.

그래서   out_features  가 1000이다.

 

만약 우리가 3개 클래스를 예측하고 싶다면   out_features  를 3으로 수정해줘야 한다.

 

수정해주는 방법은 간단하다.

수정한 structure을 그대로 재할당해주면 된다.

model_loaded.classifier = torch.nn.Sequential(
    nn.Dropout(0.2, inplace=True), 
    nn.Linear(in_features=1280,
              out_features=3)
)

print(model_loaded.classifier)
#output 
'''
Sequential(
  (0): Dropout(p=0.2, inplace=True)
  (1): Linear(in_features=1280, out_features=3, bias=True)
)

'''

 

모델 transforms확인

  weights.transforms()  는 해당 weights가 어떤 이미지 변형 방법이 적용되었는지 보여준다.

 

밑바닥부터 모델을 만들 때는   transform.compose   부터 일일이 생성해줬다.

하지만 pretrained 모델은 이미 이 부분을 구현해 주었다.

weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
transform_efficientNet = weights.transforms()

print(transform_efficientNet)

#output
'''
ImageClassification(
    crop_size=[224]
    resize_size=[256]
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    interpolation=InterpolationMode.BICUBIC
)
'''

 

실제 어떤 이미지 변형이 설정되었는지 눈으로 확인해 보자.

(transform이 낯선분은 해당 포스팅 참조)

 

train_images = glob.glob(file_path) # 전체이미지 추출

random_img_3 = random.sample(train_images, k=3) # 전체 이미지중 3개를 랜덤하게 추출

for rand_img in random_img_3:
    with Image.open(rand_img) as img:
        
        fig, ax = plt.subplots(1, 2)
        
        ax[0].imshow(img) 
        ax[0].set_title("Original \nSize: {}".format(img.size))
        ax[0].axis("off")


        img_transformed = transform_efficientNet(img).permute(1, 2, 0) 
        # shape을 [ Channels, Height, Width ] => [ Height, Width, Channels ] 변경
				# why? matplotlib에서 제대로 출력하기 위해서.

        ax[1].imshow(img_transformed) 
        ax[1].set_title("Transformed \nSize: {}".format(img_transformed.shape))
        ax[1].axis("off")

output:

 

 

이 글과 읽으면 좋은글

• pre-trained 모델 공식문서
• [pytorch] model.eval() vs torch.no_grad() 차이
• [pytorch] ImageFolder 사용 방법