LoginSignup
3
5

More than 3 years have passed since last update.

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

Last updated at Posted at 2021-04-10

前の記事で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

3
5
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
3
5