본문 바로가기

컴퓨터/머신러닝 (Machine Learning)

Pytorch, grad-cam 사용 정리

Image 분류 모델의 gradcam 사용 정리

 

pytorch grad cam 공식 깃헙

GitHub - jacobgil/pytorch-grad-cam: Advanced AI Explainability for computer vision. Support for CNNs, Vision Transformers, Classification, Object detection, Segmentation, Image similarity and more.

 

GitHub - jacobgil/pytorch-grad-cam: Advanced AI Explainability for computer vision. Support for CNNs, Vision Transformers, Clas

Advanced AI Explainability for computer vision. Support for CNNs, Vision Transformers, Classification, Object detection, Segmentation, Image similarity and more. - GitHub - jacobgil/pytorch-grad-c...

github.com

 

 

1. 패키지 설치

 

pip install grad-cam

 

2. Package import 및 함수 정의

공식 깃헙에
코드 첫줄에서 import 한 GradCAM과 AblationCAM 외에도 여러 종류의 CAM을 확인/사용할 수 있다.

 

from pytorch_grad_cam import GradCAM, AblationCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
from typing import List, Dict
import numpy as np
import cv2
import torch
from torch import nn
from torchvision.transforms import Compose, Normalize, ToTensor
import timm
import os
import time
from PIL import Image

# 공식 깃헙에서 가져옴
def preprocess_image(
    img: np.ndarray, mean=[
        0.5, 0.5, 0.5], std=[
            0.5, 0.5, 0.5]) -> torch.Tensor:
    preprocessing = Compose([
        ToTensor(),
        Normalize(mean=mean, std=std)
    ])
    return preprocessing(img.copy()).unsqueeze(0)

 

3. Pretrained 된 model load

tiny-imagenet-200 데이터 셋에서 pretrained 된 efficientnet을 불러왔다.

 

model =timm.create_model('efficientnet_lite0', pretrained=True, num_classes=200)
model.load_state_dict(torch.load('trained_imagenet.pt'))
model.to('cuda')

 

4. CAM으로 확인할 layer 확인 및 정의

CAM으로 다양한 layer를 확인할 수 있는데, batch_normalization, convolution 등이 사용되는 듯 했다.

여기서는 가장 마지막 convolution layer를 cam으로 확인하였다.

 

4-1. Layer 확인

model을 로드 한 뒤, 그냥 model을 입력하면 model의 구조가 프린트 된다. (모델 마다 다르다)

 

model

 

output:

 

 

위 사진은 프린트된 모델의 마지막 부분이다. (마지막 classifier)

사진에서 보이듯 efficientnet의 가장 마지막 convolution layer의 이름은 conv_head이다.

 

4-2. Layer 정의

 

layer = model.conv_head

 

5. CAM을 확인할 이미지 load 및 전처리

전처리는 모델을 추론할 때 데이터에 적용하는 전처리를 그대로 해줘야 한다.

(아마 대부분이 추론할 때는 resizing, normalization, reshape 쯤 할 것이다)

 

rgb_img = cv2.imread('goldfish.jpg', 1)[:, :, ::-1] #shape: (2112, 2816, 3)
rgb_img = cv2.resize(rgb_img, (342, 342))
rgb_img = np.float32(rgb_img) / 255 #shape: (342, 342, 3)

input_tensor = preprocess_image(rgb_img,
                                mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225])
# input_tensor.shape: torch.Size([1, 3, 342, 342])

 

6. target (label) 정의 및 CAM 생성

 

# target 정의, load한 이미지(금붕어)의 label이 22라서 22 입력
targets = [ClassifierOutputTarget(22)]

# cam 생성 및 처리
cam = GradCAM(model=model, target_layers=[layer], use_cuda='use_cuda')
cam.batch_size = 1
grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
grayscale_cam = grayscale_cam[0, :]

 

7. 시각화

 

visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
display(Image.fromarray(visualization))

 

output:

 

 

GradCAM의 경우 이상한 영역을 표시하였다.

하지만 AblationCAM을 사용하였을때는 아래와 같이 금붕어를 확실히 본 것으로 확인되었다.