前の記事で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に変更する必要があります。
weights_header = np.ndarray(
shape=(4, ), dtype='int32', buffer=weights_file.read(20))
変換
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
ファイルが置いてあるディレクトリになります。
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