TorchVision의 이미지 분류 모델

소개

torchvision.models 패키지는 이미지 분류, 인스턴스 분할, 물체 탐지, 인물 자세 감지, 비디오 분류 등의 다양한 작업을 수행하기 위해 사전에 학습된 모델들을 제공합니다.

이 튜토리얼에서 우리는 이미지 분류 작업에서 널리 사용되는 ResNet50 모델을 사용합니다. ResNet50 모델은 CNN 모델의 일종으로, Residual Block 구조로 연결된 50개의 레이어로 구성되어있습니다. 보다 자세한 설명은 공식 페이지에서 확인하실 수 있습니다.

모델 준비

TorchVision은 torchvision.models를 통해 여러 딥러닝 모델의 기본 구조와 그에 대응하는 가중치를 제공합니다. 이 과정은 모델 인스턴스를 초기화하는 과정에서 사전학습된 모델을 캐시 위치에 다운로드하여 불러오는 방식으로 이루어집니다.

우리는 이렇게 준비된 모델을 PyTorch에서 제공하는 코드를 사용하여 ONNX 형식으로 변환하여 사용하는 것을 추천합니다. 아래 내용을 포함한 실행가능한 파이썬 코드(예: prepare_resnet.py)를 만들어서 모델을 변환할 수 있습니다.

import torch
from torchvision.models import resnet50, ResNet50_Weights

# Using pretrained weights:
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
model.eval()

# make dummy input depending on the model's input shape
input = torch.randn(1, 3, 224, 224)

# export to onnx
torch.onnx.export(model, input, "resnet50.onnx")

Calibration Dataset 준비

TorchVision은 자신들이 모델을 학습시킨 방법에 대한 세부사항들을 제공합니다. 이러한 학습방법에는 보통 모델을 불러오는 방법, 보간법을 사용하여 이미지 크기를 재조정하는 방법, 이미지를 잘라내는 크기, 픽셀값의 조정과 정규화 등이 포함되어있습니다. 따라서 Calibration Dataset을 준비하기에 앞서 전처리에 포함된 이러한 연산들과 모델의 입력값 크기를 확인해야 합니다.

예를 들어, 우리가 사용할 ResNet50의 경우 PIL.Images 함수를 이용하여 원본 이미지를 불러오고, 이중선형보간법을 사용하여 이미지를 종횡비를 유지한 채, 짧은 변이 256이 되도록 리사이즈 하며, 중심의 (224, 224) 크기의 정사각형 부분을 잘라내고, 픽셀값을 [0, 1] 범위로 조정한 후, 평균 = [0.485, 0.456, 0.406], 표준편차 = [0.229, 0.224, 0.225]를 사용하여 정규화합니다.

ResNet50과 같은 이미지 분류 모델을 학습할 때는 보통 ImageNet 데이터셋이 사용됩니다. 따라서 이미지 종류 및 분포의 측면에서 학습에 사용한 데이터와 유사한 데이터가 Calibration Dataset으로 사용되는 것이 장려됩니다. 이 튜토리얼을 위해 모빌린트의 엔지니어들이 ImageNet의 각 클래스별로 1개의 이미지를 추출하여 두었으며, cali_1000.zip 이라는 파일명으로 제공되니, 해당 파일을 사용하시길 권장합니다.

빠른 시작에 나온 바와 같이 qubee 컴파일러는 TorchVision 패키지를 사용한 표준화된 전처리연산을 제공합니다. 따라서 Calibration Dataset은 두 가지 방법으로 만들어질 수 있습니다.

첫 번째 방법은 다음과 같은 내용을 포함한 환경설정용 YAML 파일(예: resnet50.yaml)을 이용한 것입니다.

# resnet50.yaml
Datatype: Image
GetImage:
    to_float32: false
    channel_order: RGB

Pre-Order: [ResizeTorch, CenterCrop, Normalize, SetOrder]
Pre-processing:
    ResizeTorch:
        size: 256
        interpolation: bilinear
    CenterCrop:
        size: [224, 224]
    Normalize:
        mean: [0.485, 0.456, 0.406]
        std: [0.229, 0.224, 0.225]
        to_float_div255: true
    SetOrder:
        shape: HWC

그 다음, 위 YAML 파일을 읽어서 실제 전처리를 수행하는 Python 실행 스크립트(예: prepare_calib.py)를 작성합니다.

from qubee.calibration import make_calib
make_calib(
    args_pre="/workspace/resnet50.yaml", # path to pre-processing configuration yaml file
    data_dir="/workspace/calibration/cali_1000", # path to folder of original calibration data files such as images
    save_dir="/workspace/calibration/", # path to folder to save pre-processed calibration data files
    save_name="resnet50_cali", # tag for the generated calibration dataset
    max_size=100 # Maximum number of data to use for calibration
)

두 번째 방법은 사용자 정의 전처리함수를 이용하는 방법입니다. qubee 컴파일러는 사용자가 직접 만든 전처리함수로도 손쉽게 Calibration Dataset을 구성할 수 있도록 해줍니다.

import torch
import numpy as np
from PIL import Image
import torchvision.transforms.functional as F
from torchvision.transforms import InterpolationMode
from qubee.calibration import make_calib_man

def preprocess_resnet50(img_path: str):
    img = Image.open(img_path)
    resize_size=256
    crop_size=(224, 224)
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    out = F.pil_to_tensor(img)
    out = F.resize(out, size=resize_size, interpolation = InterpolationMode.BILINEAR)
    out = F.center_crop(out, output_size=crop_size)
    out = out.to(torch.float, copy=False) / 255.
    out = F.normalize(out, mean, std)
    out = np.transpose(out.numpy(), axes=[1, 2, 0])
    return out

make_calib_man(
    pre_ftn=preprocess_resnet50, # callable function to pre-process the calibration data
    data_dir="/workspace/calibration/cali_1000", # path to folder of original calibration data files such as images
    save_dir="/workspace/calibration/", # path to folder to save pre-processed calibration data files
    save_name="resnet50_cali", # tag for the generated calibration dataset
    max_size=100 # Maximum number of data to use for calibration
)

모델 컴파일

변환된 모델과 Calibration Dataset이 준비되었다면 이제 모델을 컴파일할 수 있습니다. 예시에 사용된 ResNet50뿐만 아니라 TorchVision에서 제공되는 대부분의 이미지 분류 모델은 모빌린트의 컴파일러에서 100% 지원되며, 별도의 후처리 없이 작업에 사용할 수 있습니다.

""" Compile resnet """ 
from qubee import mxq_compile
onnx_model_path = "/workspace/resnet50.onnx"
calib_data_path = "/workspace/calibration/resnet50_cali"

mxq_compile(
    model=onnx_model_path,
    calib_data_path=calib_data_path,
    quantization_method=1, # per channel quantization
    quantization_mode=2, # max percentile quantization 
    percentile=0.999,
    quantization_output=0, # per layer quantization for the output layer
    save_path="resnet50.mxq",
    backend="onnx" 
)