Whisperモデルの軽量化 - TadaoYamaokaの開発日記

ONNXにしたWhisperのモデルの推論を速くするために、以下の方法を試した。

OpenVINOで推論
FP16化
INT8量子化

Whisperのモデルはbaseを使用している。

OpenVINOで推論

OpenVINOは、intelのCPUに最適化した推論エンジンである。
ONNX Runtimeの代わりに、OpenVINOで推論することで速くなるか試した。
ただし、実験しているCPUがRyzen Threadripperのため性能が出せない可能性もある。

インストール

はじめWindow環境にインストールしたところ、

ImportError: DLL load failed while importing _pyopenvino: 指定されたモジュールが見つかりません。

というエラーがでて実行できなかった。
そのため、WSL2のUbuntu22.04にインストールした。

OpenVINOが現時点のpytorchの最新版(1.13.1)に対応していなかったため、PyTorchのバージョンを指定してインストールする。
速度比較のためonnxruntimeもインストールする。

conda create -n openvino python=3.9.12 numpy=1.21.5
conda activate openvino
conda install pytorch=1.13.0 torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
pip install openvino-dev[onnx,pytorch] onnxruntime-gpu numpy==1.21.5 torch==1.13.0

OpenVINOで推論

OpenVINOでは以下のように推論を行う。

from openvino.runtime import Core

ie = Core()
model_encoder = ie.compile_model(model=ie.read_model(model='encoder.onnx'), device_name="CPU")

output_n_layer_cross_k = model_encoder.output(0)
output_n_layer_cross_v = model_encoder.output(1)

output_encoder = model_encoder(mel.numpy())
n_layer_cross_k = output_encoder[output_n_layer_cross_k]
n_layer_cross_v = output_encoder[output_n_layer_cross_v]

測定結果

エンコーダで1回、デコーダで2回推論を行った際の時間を計測した。
比較のため、ONNX RuntimeのCPUプロバイダでも同様に推論を行った。
3回測定して平均を算出した。

	エンコーダ	デコーダ1回目	デコーダ2回目
ONNX Runtime	0.179	0.035	0.033
OpenVINO	0.206	0.150	0.067

単位：秒

ONNX RuntimeのCPUプロバイダの方が速かった。
また、OpenVINOは、1回目の推論が遅い。

CPUがRyzen Threadripperであるためかもしれないが、少なくともAMDのCPUでも動かす場合はOpenVINOを使う必要はなさそうである。

FP16化

ONNXモデルをFP16にして速くなるか試した。

ONNXモデルのFP16化

onnxmltoolsを使ってFP16化した。

import onnx
from onnxmltools.utils.float16_converter import convert_float_to_float16_model_path

new_onnx_model = convert_float_to_float16_model_path('encoder.onnx')
onnx.save(new_onnx_model, 'encoder.fp16.onnx')

測定結果

	エンコーダ	デコーダ1回目	デコーダ2回目
ONNX Runtime(FP32)	0.179	0.035	0.033
OpenVINO(FP32)	0.206	0.150	0.067
ONNX Runtime(FP16)	0.388	0.143	0.168
OpenVINO(FP16)	0.228	0.151	0.077

ONNX RuntimeとOpenVINOのどちらもFP16化すると遅くなった。
特にONNX Runtimeは、デコーダの推論が5倍ほど遅くなる。

INT8量子化

INT8量子化は、入力値の範囲の統計を取るために、キャリブレーションが必要になるが、データセットがないため、Dynamic Quantizationを試した。

Dynamic Quantization

ONNX Runtimeを使用して、Dynamic Quantizationを以下のようにして行う。
weight_typeは、int8だと推論時にエラーが発生したため、uint8としている。

from onnxruntime.quantization import quantize_dynamic, QuantType

quantized_model = quantize_dynamic(
    'encoder.onnx',
    'encoder.int8.onnx',
    weight_type=QuantType.QUInt8,
)

測定結果

	エンコーダ	デコーダ1回目	デコーダ2回目
ONNX Runtime(FP32)	0.179	0.035	0.033
OpenVINO(FP32)	0.206	0.150	0.067
ONNX Runtime(INT8)	0.201	0.031	0.032
OpenVINO(INT8)	0.380	0.152	0.080

ONNX RuntimeはFP32とほぼ同じで、OpenVINOは少し遅くなっている。
Dynamic Quantizationによって、大幅に速くなることはなかった。

モデルサイズ

モデルサイズは以下のようになる。

	エンコーダ	デコーダ
FP32	92	296
FP16	47	152
INT8	26	77

単位：MB

INT8化によって、ファイルサイズが1/4くらいになる。

まとめ

Whisperモデルの推論を速くする方法をいくつか試した。

ONNX Runtimeの代わりにOpenVINOを使うとかえって遅くなることがわかった。
ONNX RuntimeのCPUプロバイダの方が速いというのは想定外の結果だった。

また、モデルをFP16にするとONNX Runtimeでは5倍近く遅くなることが分かった。

Dynamic Quantizationで、INT8に量子化した場合、推論速度はFP32と同じくらいであった。
INT8にするとモデルのファイルサイズが1/4くらいになるので、モデルサイズを小さくしたい場合は量子化の効果はありそうだが、推論速度はほとんど変わらないため、実行環境がPCであれば量子化しなくてもよさそうである。