YOLOv2の学習結果(weights)をonnx形式に変換する


前の記事でtiny-YOLOv2の独自データセットでの訓練を行いましたが、
このモデルをBarracudaで使用するためにはこの形式(weights/cfg)では対応できないため、
ONNX形式に変換する必要がある。
darknetで学習した成果物をONNXに変換する方法がなかなか見つからずにハマったので、整理ついでに残してみました。
他にもやり方はあると思うので、やり方の1つとして読んでいただければと思います。

変換に関する情報はそこそこあるのですが、
前回と同じく、YOLOv3以降の変換は何パターンか見つかったけど、v2についてはサポートされていないという感じでした。

darknetモデルからONNXに変換

darknetの学習結果は.weightsで出力されており、コレを最終的にはonnxにしたい。
ただ、直接ONNX形式にする方法が見つからなかったため、
一回.h5に変換して、.h5から.onnxに変換する方法を行った。
.h5はモデルがHDF5で保存されているという意味です。
kerasがこの形式の出力をサポートしているらしいです。

darknet(weights/cfg) → .h5

YAD2K: Yet Another Darknet 2 Kerasを使用して変換を行います。

YAD2K取得

$ git clone https://github.com/allanzelener/YAD2K.git

.weightsと.cfgファイルの配置

cloneしてきたYAD2Kファイル直下に以下を配置
darknetの学習結果と、学習に使ったcfgファイル

  • yolov2-tiny.weights
  • yolov2-tiny-voc.cfg

yad2k.pyのパラメータ変更

darknetを64bit端末で利用した場合、weights_file.readを16から20に変更する必要があります。

yad2k.py
    weights_header = np.ndarray(
        shape=(4, ), dtype='int32', buffer=weights_file.read(20))

参考元(日本語情報)
original issue

変換

python yad2k.py yolov2-tiny-voc.cfg yolov2-tiny.weights yolov2-tiny.h5

正常に完了すると、直下にyolov2_tiny.h5が作成されています。

.h5 → onnx

keras2onnxで変換します。
How to Convert Your Keras Model to ONNXを参考に進めました。

注意点

tensorflowのバージョンに制限があるため、
tensorflow2.4以降のversionが入っている場合は新しいpython仮想環境立てて作業することを推奨
ちなみに、tensorflow2.4の場合はエラーになります。
https://github.com/onnx/keras-onnx/issues/662

AttributeError: 'KerasTensor' object has no attribute 'graph'

必要ライブラリインストール

pip install tensorflow==2.3.0

pip install onnx

pip install keras

pip install keras2onnx

converterの準備

以下のpythonコードを準備します。
配置先は変換する.h5ファイルが置いてあるディレクトリになります。

convert_keras_to_onnx.py
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import load_model
import onnx
import keras2onnx

onnx_model_name = 'yolov2_tiny.onnx'

model = load_model('yolov2_tiny.h5')
onnx_model = keras2onnx.convert_keras(model, model.name)
onnx.save_model(onnx_model, onnx_model_name)

変換

以下コマンドでconverterを実行します

python convert_keras_to_onnx.py

成功すると、yolov2_tiny.onnxというファイルが出力されています。

参考

about barracuda

darknet → .h5

.h5 → .onnx