TorchVision (MobileNetV3 Large) -> ONNX -> TFLite (シグネチャ推論)

1. 環境

1. Intel Corei9 Gen.10
2. Ubuntu 22.04
3. TorchVision
4. ONNX
5. TensorFlow

2. パッケージのインストール

sudo apt-get update \
&& sudo apt-get install -y flatbuffers-compiler

pip install -U onnx==1.14.0 \
&& python -m pip install onnx_graphsurgeon \
    --index-url https://pypi.ngc.nvidia.com \
&& pip install -U onnx-graphsurgeon \
&& pip install -U onnxruntime==1.15.1 \
&& pip install -U onnxsim==0.4.33 \
&& pip install -U simple_onnx_processing_tools \
&& pip install -U onnx2tf \
&& pip install -U h5py==3.7.0 \
&& pip install -U psutil==5.9.5 \
&& pip install -U tensorflow==2.13.0 \
&& pip install torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cpu

3. モデル変換 → テスト推論一括実行コード

test.py

#! /usr/bin/env python

import torch
import torchvision

weights = torchvision.models.MobileNet_V3_Large_Weights.DEFAULT
model = torchvision.models.mobilenet_v3_large(weights=weights)
model.eval()
onnx_file = f'mobilenetv3_large_pytorch.onnx'
SIZE = 128
x = torch.randn((1, 3, SIZE, SIZE))
torch.onnx.export(
    model,
    args=(x),
    f=onnx_file,
    opset_version=11,
    input_names=[
        'input',
    ],
    output_names=[
        'output',
    ],
)

import onnx2tf
onnx2tf.convert(
    input_onnx_file_path=onnx_file,
    output_folder_path='saved_model',
    copy_onnx_input_output_names_to_tflite=True,
)

import time
import numpy as np
import tensorflow as tf
interpreter = tf.lite.Interpreter(
    model_path="saved_model/mobilenetv3_large_pytorch_float32.tflite"
)
tf_lite_model = interpreter.get_signature_runner()

time_total = 0.0
kaisu = 10
for i in range(kaisu):
    inputs = {
        'input': np.random.randn(1,SIZE,SIZE,3).astype(np.float32),
    }
    start_time = time.time()
    tf_lite_output = tf_lite_model(**inputs)
    elapsed_time = time.time() - start_time
    time_total += elapsed_time

print("[TFLite] Model Predictions.shape:", tf_lite_output['output'].shape)
print("[TFLite] AVG elapsed time:", time_total / kaisu)

4. 結果 (Average 6.3ms/pred)

python test.py

    :
[TFLite] Model Predictions.shape: (1, 1000)
[TFLite] AVG elapsed time: 0.006333017349243164