pretrained된 computer vision 모델에서 마지막 linear layer는 제거하고 feature만 뽑는법

예를 들어 resnet model에 대해 pretrained된 모델을 불러오고

 

import torchvision.models as models

resnet = models.resnet152(pretrained=True)

 

resnet.children()하면 resnet의 모듈을 불러올수 있다

 

 

 

마지막에 Linear() classifier layer가 있는데 얘를 제거하면 된다

 

modules = list(resnet.children())[:-1] # fully connected layer의 가장 마지막 layer 제거

 

 

이렇게 제거된 모듈을 nn.Sequential()로 sequential한 모델을 만들면 된다

 

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

resnet = models.resnet152(pretrained=True)
modules = list(resnet.children())[:-1] # fully connected layer의 가장 마지막 layer 제거
resnet = nn.Sequential(*modules)

 

 

원래 이렇게 마지막에 fc layer가 있는데

 

 

 

다음과 같이 fc layer가 없어진 sequential 모델이 나온다.

 

이는 원래 resnet의 feature extraction만 가져온 것

 

 

 

모델의 파라미터가 같아야 같은 모델인데, 실제로 같은가?

 

그러면 비교해보면 실제로 같다.

 

A = list(resnet.parameters()) #원본(467개)
B = list(resnet.parameters()) #제거(465개)

count = 0
for i in range(465):
    
    if torch.all(torch.eq(A[i],B[i])):
        
        count += 1

print(count)

 

 

보면 마지막 fc layer의 parameter는 2개이다.

 

len(list(resnet.fc.parameters()))
2

 

 

그리고 A의 마지막 두 원소 -2번, -1번이랑 비교해보면 실제로 True