Pytorch model을 ONNX로, ONNX를 openVINO로 변환하기

Intel CPU/GPU를 활용한 inference를 위해 Pytorch model을 openVINO로 변환해 본다.

- 사전 학습된 pytorch 모델이 있다고 가정.

 

참고한 링크

Convert a PyTorch Model to ONNX and OpenVINO™ IR — OpenVINO™ documentation — Version(latest)

Convert a PyTorch Model to ONNX and OpenVINO™ IR — OpenVINO™ documentation

 

1. pip 업그레이드 <- 필수

 

pip install --upgrade pip

 

2. ONNX, OpenVINO 설치

 

pip install onnx openvino

 

3. ONNX 변환

아래와 같이 Network로 정의된 모델이 있다고 할때,

 

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=5, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(12)
        self.pool = nn.MaxPool2d(2,2)
        self.conv4 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn4 = nn.BatchNorm2d(24)
        self.conv5 = nn.Conv2d(in_channels=24, out_channels=24, kernel_size=5, stride=1, padding=1)
        self.bn5 = nn.BatchNorm2d(24)
        self.fc1 = nn.Linear(24*10*10, 10)

    def forward(self, input):
        output = F.relu(self.bn1(self.conv1(input)))      
        output = F.relu(self.bn2(self.conv2(output)))     
        output = self.pool(output)                        
        output = F.relu(self.bn4(self.conv4(output)))     
        output = F.relu(self.bn5(self.conv5(output)))     
        output = output.view(-1, 24*10*10)
        output = self.fc1(output)

        return output

 

변환할 모델과 웨이트 파일을 로드한다.

 

import torch

model = Network()
model.load_state_dict(torch.load("웨이트파일 경로", map_location=None))
model.eval()

 

dummy input을 정의한다.

 

# IMAGE_HEIGHT와 IMGAE_WIDTH는 모델에 들어가는 이미지의(resize 된) 크기
# 나머지는 고정
dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)

 

ONNX 변환 코드

 

#onnx_저장경로 변경!
torch.onnx.export(
    model,
    dummy_input,
    onnx_저장경로,
    opset_version=11,
    do_constant_folding=False,
)

 

4. OpenVINO 변환

아래의 코드는 openvino 변환에 사용될 command를 만들어준다.

 

# Construct the command for Model Optimizer.
# onnx_저장경로, openvino_저장경로, IMAGE_HEIGHT, IMAGE_WIDTH 변경!
from IPython.display import Markdown, display

mo_command = f"""mo
                 --input_model "onnx_저장경로"
                 --input_shape "[1,3, {IMAGE_HEIGHT}, {IMAGE_WIDTH}]"
                 --data_type FP16
                 --output_dir "openvino_저장경로"
                 """
mo_command = " ".join(mo_command.split())
print("Model Optimizer command to convert the ONNX model to OpenVINO:")
display(Markdown(f"`{mo_command}`"))

 

output:

 

 

OpenVINO 변환코드

 

print("Exporting ONNX model to IR... This may take a few minutes.")
mo_result = %sx $mo_command
print("\n".join(mo_result))

 

output:

 

 

끝