kerasで作った学習済みモデルを.pbファイルに変換する

ソフトウェア

 

kerasで作った学習済みモデルをpbファイルに変換したので備忘録です。

 

やりたかったこと

 

やりたかったことは、

kerasで作った学習済みモデルをOpenVINOに変換

です。

 

複数の画像を同時かつ高速にAIで判定する必要があり、kerasで実現が難しかったので

CPUで高速に推論が可能なOpenVINOを採用することにしました。

 

kerasで実現しようと努力した経緯も記事にしていますので、興味がある方はぜひ。

https://nori-life-log.com/flaskでkerasの推論サーバ

 

kerasで作った学習済みモデルをOpenVINO形式に変換するには、まずTensorFlow形式、いわゆる.pbファイルに変換する必要があったので、いろいろ調べて実装しました。

 

 

コード

 

コードは以下です。

下記ブログを参考にしました。

Just a moment...
import tensorflow as tffrom keras.models import load_model
from keras import backend as K
from tensorflow.python.framework import graph_io
from tensorflow.python.framework import graph_util


def freeze_graph(session, output, save_pb_dir='.', save_pb_name='frozen_model.pb', save_pb_as_text=False):
    graph = session.graph
    with graph.as_default():
        graphdef_inf = tf.graph_util.remove_training_nodes(graph.as_graph_def())
        graphdef_frozen = graph_util.convert_variables_to_constants(session, graphdef_inf, output)
        graph_io.write_graph(graphdef_frozen, save_pb_dir, save_pb_name, as_text=save_pb_as_text)
        return graphdef_frozen

def main():
    save_pb_dir = 'your_pb_dir'
    keras_model = 'your_keras_model.h5'

    # Clear any previous session.
    K.clear_session()
    # This line must be executed before loading Keras model.
    K.set_learning_phase(0)
    model = load_model(keras_model)
    session = K.get_session()
    freeze_graph(session, [out.op.name for out in model.outputs], save_pb_dir=save_pb_dir)

if __name__ == '__main__':
    main()
 

 

参考にしたブログから変更したポイントは、

tensorflowのbackendではなくkerasのbackendを使っている点です。

 

最初は参考ブログと同様に

import tensorflow as tf

tf.keras.backend.hogehoge

と書いていたのですが、なぜかうまく動きませんでした。

 

他の記事を見ると、代わりに

from keras import backend as K

K.backend.hogehoge

と書いてある記事があったので、試してみるとうまく動きました。

 

まとめ

kerasの学習済みモデルを.pbファイルに変換した際の備忘録でした。

tensorflow.keras.backendとkeras.backendの違いが良くわからないままですが、動いたので次に進みます(笑)。

 

コメント