[CV] ResNet-18로 특정 Image의 feature 추출하기 (PyTorch)

오늘은 리부탈에서 사용했던 간단한 feature extraction 과정을 포스팅해보려고 한다. 리부탈 끝난 기념으로 여유롭게 포스팅을 하구있다 🥳🔥

우선 오늘 포스팅할 내용은 특정 Image 하나가 들어왔을 때, 이 image를 잘 나타내는 feature를 추출하는 pytorch 방법이다. 나의 경우에는 segmentation map의 feature가 필요했고, 단순 convolution layer를 쌓는 것보다 pretrain된 PyTorch 내장 모듈을 사용하는 것이 좋을 것 같다는 생각에 ResNet을 이용하게 됐다. 

우선 간단히 ResNet-18부터 짚고 넘어가고, 어떻게 특정 image의 feature를 추출하는지 그 code를 소개해보려고 한다. 

 

# ResNet-18

우선 이 글을 읽는 독자라면 ResNet은 익숙할 것이다. (만약 익숙치 않다면 여기에 아주 간단한 설명이 나와있다!) ResNet-18은 말그대로 그냥 18개의 층으로 이루어진 ResNet을 의미한다. 간단하게 ResNet의 구조에 대해 설명하면 다음과 같다. 

ResNet-18의 구조

우선 ResNet-18의 input은 244 * 244 * 3 의 image이다. 그리고 위 그림처럼 4개의 Conv block을 거치고, 마지막에 Adaptive Average Pooling 을 시행한 후 FC(Fully Connected) layer를 통과시켜 이미지 분류를 수행한다. 참고로 지금 우리는 classification이 아니라 단지 여러개의 convolution block을 거치며 나온 feature를 추출하고 싶은 것이므로 FC layer는 사용하지 않는다. 

우리는 당연히 위 ResNet을 직접 train 시키는 것이 아니라, 이미 누군가가 열심히 train 해놓은 pretrain 모델을 feature extraction에 사용할 것이다. PyTorch에서는 이러한 pretrain 모델을 다음과 같은 코드로 사용할 수 있다. 

import torchvision.models as models
model = models.resnet18(pretrained=True)

우리는 위 pretrain 모델을 불러온 후, model을 eval 모드로 바꾼 다음 feature extraction을 수행하면 된다. 

 

---

# Feature extraction 

그럼 이제부터 이 포스팅의 주제였던 특정 image의 feature를 ResNet-18로 추출하는 과정을 소개하겠다. 전반적인 코드는 여기를 참고하였다. 

 

1. Image prep

우선 앞서 소개했 듯 pretrain된 ResNet의 input size는 224 * 224 * 3으로 정해져있기 때문에, 우리의 image도 이러한 차원을 맞춰줘야 한다. 따라서 모델에 이미지를 무작정 넣기 전에 사전 준비하는 과정이 필요한데, image size를 맞춰주는 과정 외에도 정규화(Normalize) 과정이 필요하다. 이러한 과정은 PyTorch에서 다음과 같이 수행할 수 있다. 

import torchvision.transforms as transforms

scaler = transforms.Scale((224, 224))         # resize 
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])   # normalize 
to_tensor = transforms.ToTensor()   # to tensor

 

2. Model load 

그 다음은 앞서 소개한 pretrain된 Resnet model을 사용하기 위해 model을 load 해준다. 그리고 위 ResNet 구조 중 우리는 사실 여러 Conv block을 거쳐 추출된 Image feature vector를 원하는 것이므로, 모델 구조 중 우리는 마지막 부분인 "Avg pool" 에 관심있다. 따라서 Avg pool layer를 아래 코드 처럼 선택해준다. 마지막으로 우리는 ResNet-18을 직접 train할 것이 아니기때문에, model.eval()도 코드에 추가해 evaluation mode로 사용할 것임을 명시해준다! 

model = models.resnet18(pretrained=True)    # pretrain model load 
layer = model._modules.get('avgpool')       # avg pool layer 선택
model.eval()                             # 평가 모드로 바꾼다

 

3. Layer connection 

마지막으로는 추출된 feature를 담을 벡터를 생성하고, 우리가 위에서 선택한 avgpool layer에 layer의 output을 복사할 함수를 연결하는 것이다. 우선 추출된 feature를 저장할 공간의 size는 아마 avg pool의 output size와 같을 것이다. 여기서 avg pool의 layer 출력 크기는 512 이므로, 512 size의 zero vector를 임의로 생성해준다. 

또한 register_forward_hook(copy_data) 을 이용해 선택한 avg pool layer에 우리가 얻고싶은 벡터를 복사할 함수를 연결해준다. 이러한 과정을 코드로 나타내면 다음과 같다. 

my_embedding = torch.zeros(512)      # 특징 벡터를 저장할 공간 정의 

def copy_data(m, i, o):           # 특징 벡터를 복사할 copy layer 정의 
    my_embedding.copy_(o.data)

h = layer.register_forward_hook(copy_data)   # register_forward_hook을 이용해 위 copy layer 연결

 

4. Overall function 

위에서 소개한 모든 과정을 하나의 함수로 나타내면 다음과 같다. 

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
from PIL import Image

# Load the pretrained model
model = models.resnet18(pretrained=True)

# Use the model object to select the desired layer
layer = model._modules.get('avgpool')

# Set model to evaluation mode
model.eval()

# Image transforms
scaler = transforms.Scale((224, 224))
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
to_tensor = transforms.ToTensor()

def get_vector(image_name):
    # 1. Load the image with Pillow library
    img = Image.open(image_name)
    # 2. Create a PyTorch Variable with the transformed image
    t_img = Variable(normalize(to_tensor(scaler(img))).unsqueeze(0))
    # 3. Create a vector of zeros that will hold our feature vector
    #    The 'avgpool' layer has an output size of 512
    my_embedding = torch.zeros(512)
    # 4. Define a function that will copy the output of a layer
    def copy_data(m, i, o):
        my_embedding.copy_(o.data)
    # 5. Attach that function to our selected layer
    h = layer.register_forward_hook(copy_data)
    # 6. Run the model on our transformed image
    model(t_img)
    # 7. Detach our copy function from the layer
    h.remove()
    # 8. Return the feature vector
    return my_embedding.numpy()

 

지금까지 Pretrain 된 ResNet으로 image feature를 추출하는 방법에 대해 포스팅해보았다. :)
참고로 output size인 512를 변경해주고 싶다면, 나는 추출된 feature 뒤에 linear layer를 붙여 내가 필요한 차원으로 변경해주었다. 

 

---

References 

[0] https://hnsuk.tistory.com/31

[1] https://becominghuman.ai/extract-a-feature-vector-for-any-image-with-pytorch-9717561d1d4c

[2] https://github.com/josharnoldjosh/Resnet-Extract-Image-Feature-Pytorch-Python/blob/master/extract_image_feature.py