Pythonでネットワークを学習し、C++で推論する （2022年版）
TenosrFlowを用いて、PythonとGPUでネットワークのパラメータを学習し、C++とCPUで推論を行うために必要な作業を紹介します。
前回は、前準備としてネットワークの入出力層の名前を取得する方法を説明しました。
前提条件
TensorFlowを用いた学習はPythonで実行します。
モデルのパラメータやネットワークを保存した.ckptファイルはすでに作成済みであるとします。
Pythonで学習したネットワークを用いてC++で推論を行います。
また、学習・推論で用いるアクセラレータ等の環境はCPU onlyまたはCPU+GPUを想定しています。
ネットワークの入出力層の名前はすでに分かっている（詳しくはこちらをご参照ください）。
環境
学習：tensorflow-gpu-2.4.1（Python）
推論：tensorflow_cc-2.4.1（C++）
C++でTensorFlowの推論を実行する手順
[1. 学習済みモデルの入力層・出力層の名前を取得する。](https://qiita.com/kyaFUK/items/635ef0d31e8e7eba9ab8) 2. **.ckptファイルを.pbファイルに変換する。** [3. `.pb`ファイルとtensorflow_ccを用いて、`C++`で`TensorFlow`の推論を実行する。](https://qiita.com/kyaFUK/items/e4dd03301e601298504f)このページでは2.を解説します。
.ckptファイルを.pbファイルに変換する。
TensorFlowにおいてモデルの保存には、PythonではCheckpointファイル (.ckpt)、C++ではprotobufファイル (.pb)1 を用います。
以下のプログラムを用いてPythonで学習したモデルをfreeze(モデルを再training不可能にし、パラメータを固定する)し、C++で読込可能な形式に変換します。
import tensorflow.compat.v1 as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # If you want to use only CPUs
output_node_names = ['stack']
with tf.Session() as sess:
# Restore the graph
saver = tf.train.import_meta_graph("./checkpoints/model.ckpt-XXXX.meta")
# Load weights
saver.restore(sess, './checkpoints/model.ckpt-XXXX')
# Set the device for CPU
assing_ops = tf.Graph.get_operations(sess.graph)
for op in assing_ops:
op._set_device('/device:CPU:*')
# Freeze the graph
frozen_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
sess.graph_def,
output_node_names
)
# Save the frozen graph
with open('frozen_graph.pb', 'wb') as f:
f.write(frozen_graph_def.SerializeToString())
freezingにGPUを利用しないことを明示する
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # If you want to use only CPUs
前回調べた出力層の名前を明示する。
output_node_names = ['stack']
モデルのネットワークの情報と学習済みパラメータを読み込む
.ckptファイルには、ネットワークの構造と学習済みパラメータが保存されています。
.ckptファイルの構成の一例を以下に示します。
checkpoints
├─model.ckpt-XXXX.meta
├─model.ckpt-XXXX.data-00000-of-00001
└─model.ckpt-XXXX.index
model.ckpt-XXXX.metaには、ネットワークの構造が保存されており、model.ckpt-XXXX.data-00000-of-00001には、学習済みパラメータが保存されています 2。
以下のようにモデルの情報を読み込みます。
# Restore the graph
saver = tf.train.import_meta_graph("./checkpoints/model.ckpt-XXXX.meta")
# Load weights
saver.restore(sess, './checkpoints/model.ckpt-XXXX')
推論にはCPUのみ利用することを明示する。
GPUを用いて学習した場合、GPUで並列化するための情報が記録されています。
すべての計算にCPUを用いることを明示します。
この部分は、学習 (GPU) -> 推論 (CPU)の場合にのみ必要です。
# Set the device for CPU
assing_ops = tf.Graph.get_operations(sess.graph)
for op in assing_ops:
op._set_device('/device:CPU:*')
ネットワーク・パラメータの情報を固定する。
convert_variables_to_constants()を使って、パラメータを定数に変換します。
frozen_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
sess.graph_def,
output_node_names
)
次の手順
3. .pbファイルとtensorflow_ccを用いて、C++でTensorFlowの推論を実行する。
（追記しました。）
参考にしたサイト
TensorFlowの学習・推論をC++で行う。
.ckptファイルと.pbファイルの変換
.ckptファイル自体の解説