GGRS: Geoscience, GIS, & Remote Sensing

지구과학, GIS, 그리고 원격탐사 블로그입니다.

PyTorch: 사전 학습 모델을 이용한 이미지 분류

댓글 0

IT

2021. 7. 21.

안녕하세요? 이번 글은 PyTorch에서 사전 학습 모델(pre-trained model)을 이용한 이미지 분류(image classification) 방법을 정리해 보겠습니다. 이미지 분류를 위해 TorchVision 모듈에 있는 사전 훈련 모델을 사용하는 방식입니다.

 

PyTorch 설치 과정은 다음 글을 참조하실 수 있습니다.

 

PyTorch: PyTorch(파이토치) 설치하기

안녕하세요? 이번 글은 PyTorch(파이토치) 설치 과정을 정리해 보겠습니다. PyTorch는 공식 홈페이지는 다음과 같습니다: pytorch.org/ PyTorch An open source deep learning platform that provides a seamless..

blog.daum.net

 

PyTorch는 오픈소스 머신러닝 프레임워크이며 TorchVision 라이브러리는 PyTorch 프로젝트의 일부입니다. torchvision 패키지는 일반적인 데이터셋, 모델 아키텍처 및 컴퓨터 비전을 위한 범용적 이미지 변환들로 구성됩니다.

 

torchvision — Torchvision 0.10.0 documentation

Shortcuts

pytorch.org

(참고) PyTorch 설치(https://pytorch.org/get-started/locally/) 시, torchvision 패키지는 함께 설치됩니다.

일단, torchvision 모듈에서 사용가능한 모델을 확인해 봅니다.

from torchvision import models
import torch

dir(models)
['AlexNet',
 'DenseNet',
 'GoogLeNet',
 'GoogLeNetOutputs',
 'Inception3',
 'InceptionOutputs',
 'MNASNet',
 'MobileNetV2',
 'MobileNetV3',
 'ResNet',
 'ShuffleNetV2',
 'SqueezeNet',
 'VGG',
 '_GoogLeNetOutputs',
 '_InceptionOutputs',
 '__builtins__',
 '__cached__',
 '__doc__',
 '__file__',
 '__loader__',
 '__name__',
 '__package__',
 '__path__',
 '__spec__',
 '_utils',
 'alexnet',
 'densenet',
 'densenet121',
 'densenet161',
 'densenet169',
 'densenet201',
 'detection',
 'googlenet',
 'inception',
 'inception_v3',
 'mnasnet',
 'mnasnet0_5',
 'mnasnet0_75',
 'mnasnet1_0',
 'mnasnet1_3',
 'mobilenet',
 'mobilenet_v2',
 'mobilenet_v3_large',
 'mobilenet_v3_small',
 'mobilenetv2',
 'mobilenetv3',
 'quantization',
 'resnet',
 'resnet101',
 'resnet152',
 'resnet18',
 'resnet34',
 'resnet50',
 'resnext101_32x8d',
 'resnext50_32x4d',
 'segmentation',
 'shufflenet_v2_x0_5',
 'shufflenet_v2_x1_0',
 'shufflenet_v2_x1_5',
 'shufflenet_v2_x2_0',
 'shufflenetv2',
 'squeezenet',
 'squeezenet1_0',
 'squeezenet1_1',
 'utils',
 'vgg',
 'vgg11',
 'vgg11_bn',
 'vgg13',
 'vgg13_bn',
 'vgg16',
 'vgg16_bn',
 'vgg19',
 'vgg19_bn',
 'video',
 'wide_resnet101_2',
 'wide_resnet50_2']

저는 'AlexNet' 모델을 사용해 보겠습니다.

model = models.alexnet(pretrained=True)
model.eval()
AlexNet(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): ReLU(inplace=True)
    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (9): ReLU(inplace=True)
    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))
  (classifier): Sequential(
    (0): Dropout(p=0.5, inplace=False)
    (1): Linear(in_features=9216, out_features=4096, bias=True)
    (2): ReLU(inplace=True)
    (3): Dropout(p=0.5, inplace=False)
    (4): Linear(in_features=4096, out_features=4096, bias=True)
    (5): ReLU(inplace=True)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

실습용으로 아래 이미지(출처: 미국 국립공원청)를 사용해 봤습니다.

import urllib
url, filename = ("https://www.nps.gov/grsm/learn/nature/images/big-bear-3-carter_1.jpg", "bear.jpg")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
input_image = Image.open(filename)
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

# move the input and model to GPU for speed if available
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    model.to('cuda')

with torch.no_grad():
    output = model(input_batch)
# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
print(output[0])
# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
probabilities = torch.nn.functional.softmax(output[0], dim=0)
print(probabilities)
# Download ImageNet labels
!wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt
# Read the categories
with open("imagenet_classes.txt", "r") as f:
    categories = [s.strip() for s in f.readlines()]
# Show top categories per image
top5_prob, top5_catid = torch.topk(probabilities, 5)
for i in range(top5_prob.size(0)):
    print(categories[top5_catid[i]], top5_prob[i].item())

위 코드는 아래 링크의 'AlexNet | PyTorch'에서 제공하는 소스를 그대로 사용했습니다.

 

PyTorch

An open source machine learning framework that accelerates the path from research prototyping to production deployment.

pytorch.org

결과는 다음과 같습니다. 아메리카흑곰(American black bear) 이미지로 분류되었네요!

American black bear 0.49522098898887634
sloth bear 0.2717072665691376
schipperke 0.13724419474601746
brown bear 0.02996181510388851
groenendael 0.01665865071117878