はじめに
現在主流のDeep Learningフレームワークとしては
- Tensorflow2
- PyTorch
- Keras
- etc...
といろいろありますが、作成した学習済みモデルの保存形式は以下表のとおりフレームワークごとに異なります。
| フレームワーク | 保存形式 | 備考 |
|---|---|---|
| Tensorflow2 | .pb | tf2onnxを使うことでonnx形式も可能っぽい |
| PyTorch | .pt, .onnx | - |
| Keras | .h5 | keras2onnxを使うことでonnx形式も可能 |
私個人としては以下の悩みがあり、
- 個人的にKerasに慣れているのでKerasで学習したいけど、保存形式は.pbにしたい。
- .pbで保存する方法をネットで検索しても出てくる方法はTensorflowがver1の時の方法なので使えない。
上で挙げた私の悩みを皆さんもお持ちのようなので、現状私の環境で成功したコードを置いておきます。
ただ、残念なことに私自身Tensorflowに詳しくないため、質問をしていただいてもたぶん適切な回答は出来ないと思います...
環境
このエントリを書いている私の環境は以下です。
- OS : Windows10 Pro x64 20H2
- CPU : Intel
- Python3 : 3.7.9
- Keras : 2.4.3
- tendorflow : 2.3.0
- tensorflow-gpu : 2.3.0
コード
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
#path of the directory where you want to save your model
frozen_out_path = ''
# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()
layers = [op.name for op in frozen_func.graph.get_operations()]
# Save frozen graph to disk
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir=frozen_out_path,
name="mnist.pb",
as_text=False)
# Save its text representation
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir=frozen_out_path,
name="mnist.pbtxt",
as_text=True)
おわりに
私が.pbで保存したい理由は 学習済みモデルをOpenVINOで使いたかったからです。
そちらはもうそろそろ記事にしますので、この記事と一緒に見ていただけると嬉しいです。
誤字、脱字、認識間違いなどあればご指摘をお願い致します。