OpenVINO – PyTorch 모델을 ONNX/OpenVINO IR 모델로 변환하기

딥러닝 모델을 학습하고 배포하기 위해서는 모델을 효율적으로 변환하고 최적화하는 과정이 필요합니다. PyTorch로 학습한 모델을 OpenVINO에서 사용하기 위해서는 ONNX 형식으로 변환한 뒤 OpenVINO Intermediate Representation(IR) 모델로 변환해야 합니다.

아래는 PyTorch 모델을 ONNX 형식으로 변환하는 코드입니다:

import torch
import torchvision
import onnx

# PyTorch 모델 불러오기
model = torchvision.models.resnet50()
model.load_state_dict(torch.load('resnet50.pth'))

# 입력 데이터 생성
dummy_input = torch.randn(1, 3, 224, 224)

# ONNX 모델로 변환
torch.onnx.export(model, dummy_input, 'resnet50.onnx')

이제 ONNX 형식으로 저장된 모델을 OpenVINO IR 모델로 변환할 차례입니다. 아래는 변환하는 코드입니다:

import onnx
from openvino.inference_engine import IECore

# ONNX 모델 로드
model = onnx.load('resnet50.onnx')

# OpenVINO IR 모델로 변환
ie = IECore()
net = ie.read_network(model, 'IR')
net.save_model('resnet50.xml')

위의 코드를 실행하면 PyTorch로 학습한 모델을 OpenVINO에서 사용할 수 있는 형식으로 변환할 수 있습니다. 이를 통해 모델을 효율적으로 배포하고 최적화할 수 있습니다.