Pytorch로 학습한 Image AI모델을 Torchscript로 변환 및 실행방법

 

 Torchscript에 대해 검색을 하는 과정에서 제 글을 찾기까지 여러분께서는 Torchscript에 대해 많은 내용들을 배우셨으리라 생각합니다.
  그럼에도 수많은 Torchscript에 다룬 글들을 끊임없이 찾다 제 블로그를 발견하신 여러분들이라면 아마도 이러한 고민을 하셨을 것으로 어림 짐작해봅니다.

 

"간단한 예제들을 보면서 Torchscript에 대해 이해할 수 있었습니다. 하지만 내가 학습한 AI 모델을 Torchscript로 변환하려면 어떻게 해야하죠?

 

 저 또한 지금까지 Torchscript 이론 및 예제들을 익혀보았지만, 실무에 바로 적용하기엔  뭔가 부족한 느낌을 많이 받았습니다.  카메라 영상 데이터를 사용하여 AI모델을 만들고 있는 저에게 있어 영상 데이터를 입력값으로 하는 Torchscript 예제를 찾기가 상당히 어려웠습니다. 다행히도 수없이 많은 구글링과 ChatGPT의 도움 덕에 Torchscript모델을 현업에 적용하는데 간신히 성공하였습니다.

 

  이 포스팅에서는 컬러 이미지를 입력으로 받는 AI모델을 Torchscript로 변환한 후 활용하는 과정들을 기록하였습니다. 혹시나 Torchscript에 대해 이론적인 지식을 필요로 하시는 분들께서는 아래의 자료를 참조해주셨으면 합니다.

https://velog.io/@dev_junseok_22/TorchScript-%EC%86%8C%EA%B0%9C

TorchScript 소개

 

아래 코드를 통해 Vision AI모델을 TorchScript로 변환하고 실행하는 예제를 확인하실 수 있습니다.

 

먼저, Pytorch로 학습한 AI모델을 TorchScript로 변환합니다 여기서는 TorchVision에서 제공되는 Pretrained Model을 변환해보겠습니다.

import torch
import torchvision
 
# TorchVision에서 제공하는 모델을 로드
model = torchvision.models.resnet18(pretrained=True)
# 모델을 평가 모드로 전환
model.eval()
# Trace 방식으로 TorchScript 변환
example_input - torch.randn(1, 3, 512, 512)
traced_model = torch.jit.trace(model, example_input)
# TorchScript 모델 저장
traced_model.save('end2end.pt')

 

다음은 TorchScript로 변환된 AI모델을 Inference 하는 방법입니다.

 

import cv2
import numpy as np
import torch
import os
import glob
from torchvision import transforms
 
# AI모델을 학습하였을 때 사용한 이미지 Dataset의 정보
mean=[95.5142,114.2475,112.4564]
std=[40.0954,43.8636,38.223]
input_size=(512,512)
 
# 입력 데이터로 사용할 이미지 불러오기
image_path = "input_image.jpg"
img = cv2.imread(image_path)
height, width, c = img.shape
 
# AI모델의 Input 양식으로 변환
inputs = cv2.resize(img, input_size)
inputs = (inputs - mean) / std
 
# 이미지 데이터를 PyTorch 텐서로 변환
toTensor = transforms.ToTensor()
input_tensor = toTensor(inputs.astype(np.float32))
 
# 입력 이미지의 Tensor 차원을 추가함
# RGB 이미지 기준으로 Tensor가 (3,512,512)에서 (1,3,512,512)로 차원이 증가함
# Batch Size가 1인 입력으로 만들어 AI모델에 입력하는 것임
# to('cuda')로 Tensor값을 GPU메모리에 올림
input_batch = input_tensor.unsqueeze(0).to('cuda:1')
 
# 저장된 TorchScript 모델 로드
script_model = torch.jit.load('end2end.pt')
# 모델을 GPU로 이동
script_model.to('cuda:1')
# 모델을 추론 모드로 설정
script_model.eval()
 
# 모델 추론
with torch.no_grad():
    output = script_model(input_batch)
 
# 추론 결과를 GPU에서 CPU 메모리로 불러옴
inference_result = output.squeeze().cpu().numpy()
 
# 추론 결과를 이미지의 원본 크기로 Resize하여 결과 출력
seg = cv2.resize(inference_result, (width, height), interpolation=cv2.INTER_NEAREST)
 

 

이번 포스팅을 작성하면서 딥러닝 AI모델을 가공하는 과정들에 대해 다시 한 번 복습할 수 있었습니다.

실무자 입장에서 오픈소스로 공개된 라이브러리들이 이미지 전처리를 포함한 전반적인 학습 과정들을 포함해 잘 갖추어져 있어 가공되지 않은 원본 이미지를 입력으로 하면 곧바로 결과를 확인할 수 있다보니 Vision AI를 연구하는 입장으로서 AI모델의 학습과정에 필요한 기초적인 내용들에 대해 어렴풋이만 알고 있었다가 Torchscript로 변환된 모델을 다루게 되면서 Low 데이터로부터 학습데이터를 가공하는 기술의 중요성을 깨닫게 되었습니다.

 

 최신 AI 기술들의 경향을 따라가는 것은 중요하지만, 가끔은 AI의 기초를 다시 한 번 되짚어보는 것도 좋을것같습니다.