🍷

OpenVINOで各フレームワーク（ONNX、TensorFlow、TFLite、PaddlePaddle）の重みを読み込んで推論

2024/08/10に公開

ただし、PyTorchさん、あなたはダメです。直接読み込めません👻

今回のサンプルは、Google Colaboratory上でお試しします。
いつの間にか、Colaboratory上でのOpenVINOインストール・実行も簡単になってて、助かります。
個人的には、PaddlePaddleの重みを直接読み込めるのが良い感じです。

ノートブックは以下のリポジトリで公開しています。
試してみたい方は「Open in Colab」からノートブックを開いて、上から順に実行していってください。

以降は処理の簡単な説明です。

パッケージインストール

OpenVINOのインストールです。
今回のサンプルを動かすだけであれば、pipインストール1行で済みます。

!pip install -q "openvino>=2024.3.0"

今回のサンプルでは、PyTorchからONNXモデルへ変換し、ONNXから各フレームワークの重みへ変換します。
それらの実行に必要なパッケージをインストールします。
ONNXからTensorFlow、TensorFlow Liteへの変換は、PINTOさんのonnx2tfをお借りしました。ありがとうございます🦔

!pip install -q onnx
!pip install -q onnx-graphsurgeon sng4onnx onnxsim onnx2tf

!pip install -q paddlepaddle x2paddle

モデルファイル準備

今回は、EfficientNet-B0を使用しました。

まずは、TorchvisionのmodelsからEfficientNet-B0を読み込み、
PyTorchの重み(*.pth)を保存します。

import torch
from torchvision import models
from torchvision.models import EfficientNet_B0_Weights

# PyTorch
pytorch_model = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)
pytorch_model.eval()
torch.save(pytorch_model, 'efficientnet_b0.pth')

次にPyTorchからONNXの重み(*.onnx)をエクスポートします。

# PyTorch -> ONNX
torch.onnx.export(
    pytorch_model,
    torch.randn(1, 3, 224, 224),
    'efficientnet_b0.onnx',
    input_names=['input'],
    output_names=['output'],
)

ONNXからOpenVINOの重み(*.xml、*.bin)に変換して保存します。

import openvino as ov

# ONNX -> OpenVINO
ov_model = ov.convert_model('efficientnet_b0.onnx')
ov.save_model(ov_model, 'efficientnet_b0.xml')

onn2tfを用いて、ONNXからTensorFlowの重み(*.pb、*.tflite)に変換します。

# ONNX -> TensorFlow, TensorFlow Lite
!onnx2tf -i efficientnet_b0.onnx -otfv1pb
%cp saved_model/efficientnet_b0_float32.pb ./efficientnet_b0.pb
%cp saved_model/efficientnet_b0_float32.tflite ./efficientnet_b0.tflite

x2paddleを用いて、ONNXからPaddlePaddleの重み(*.pdmodel、*.pdiparams)に変換します。

# ONNX -> PaddlePaddle
!x2paddle --framework=onnx --model=efficientnet_b0.onnx --save_dir='pd_model'
%cp pd_model/inference_model/model.pdmodel ./efficientnet_b0.pdmodel
%cp pd_model/inference_model/model.pdiparams ./efficientnet_b0.pdiparams

サンプル画像取得

サンプル画像を取得します。
サンプル画像はぱくたそ様の「ハーブから吸蜜する働き蜂」を使用しています。

!pip install -q imread_from_url

from imread_from_url import imread_from_url

image = imread_from_url('https://raw.githubusercontent.com/Kazuhito00/OpenVINO-Colab-Inference-Sample-using-various-frameworks-weights/main/sample.jpg')

from google.colab.patches import cv2_imshow

cv2_imshow(image)

読み込んだ画像に対し前処理（RGB変換、リサイズ、正規化、BCHW変換）を実施します。
この後の推論で、TensorFlow系のモデルはbhwc_image、それ以外はbchw_imageを使用します。

import cv2
import numpy as np

def preprocess_image(image):
    preprocessed_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    preprocessed_image = cv2.resize(preprocessed_image, (224, 224))

    preprocessed_image = preprocessed_image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    preprocessed_image = (preprocessed_image - mean) / std

    bhwc_image = np.expand_dims(preprocessed_image, axis=0)

    bchw_image = np.transpose(preprocessed_image, (2, 0, 1))
    bchw_image = np.expand_dims(bchw_image, axis=0)

    return bchw_image, bhwc_image

# 前処理
bchw_image, bhwc_image = preprocess_image(image)

PyTorch推論テスト

まずは、OpenVINOで推論する前に、PyTorchで推論して動作確認します。
推論して、上位3の結果を表示します。

image_tensor = torch.tensor(bchw_image, dtype=torch.float32)

# 推論
with torch.no_grad():
    output = pytorch_model(image_tensor)

# 推論結果確認
result = output.numpy().squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

それぞれのラベルは以下なので、あっていそうですね👀
309：Bee(蜂)
94：hummingbird(ハチドリ)
308：fly(蝿)

OpenVINO推論テスト

OpenVINOをインポートして、OpenVINOコアのインスタンスを生成します。

import openvino as ov

core = ov.Core()

OpenVINO(.xml、.bin)

まずは、OpenVINOの重みを読み込んで推論します。
xmlを読み込むと、binも読み込まれます。

# OpenVINOモデル読み込み
ov_model = core.read_model(model='efficientnet_b0.xml')
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bchw_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

ONNX(*.onnx)

次に、ONNXの重みを読み込んで推論します。
重みのパス指定以外は、全て共通のAPIコールで動作します。
以降の他フレームワークの重みも同様です。

# ONNXモデル読み込み
ov_model = core.read_model(model='efficientnet_b0.onnx')
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bchw_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

TensorFlow(*.pb)

TensorFlowの重みを読み込んで推論します。
TensorFlow系だけは入力がBHWCです。

# TensorFlowモデル読み込み
ov_model = core.read_model(model='efficientnet_b0.pb')
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bhwc_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

TensorFlow Lite(*.tflite)

TensorFlow Liteの重みを読み込んで推論します。

# TensorFlow Liteモデル読み込み
ov_model = core.read_model(model='efficientnet_b0.tflite')
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bhwc_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

PaddlePaddle(.pdmodel、.pdiparams)

PaddlePaddleの重みを読み込んで推論します。

# PaddlePaddleモデル読み込み
ov_model = core.read_model(model='efficientnet_b0.pdmodel')
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bchw_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)

それぞれ、若干のスコア差異はありますが、推論結果は「309：Bee(蜂)」となっているので、問題無く使用できそうです。

PyTorchの重みを使用してOpenVINOで推論

PyTorchの重みはOpenVINOで直接読み込めないため、convert_model()で読み込んで使用します。
と言うか、これまんまOpenVINOへの変換処理なので、つまりOpenVINOの重みで推論しているのと同一ですが。。。

# PyTorch(直接読み込めないため、convert_model()使用)
ov_model = ov.convert_model(pytorch_model, example_input=torch.zeros((1, 3, 224, 224)))
ov_compile_model = core.compile_model(model=ov_model, device_name="CPU")

# 推論
output = ov_compile_model([bchw_image])
output_layer = ov_compile_model.output(0)
result = output[output_layer]

# 推論結果確認
result = result.squeeze()
top3_index = result.argsort()[-3:][::-1]
top3_score = result[top3_index]
print(top3_index)
print(top3_score)