Pythonでネットワークを学習し、C++で推論する （2022年版）
TensorFlowを用いてDeep Learningを行う場合、Pythonを使うのが一般的になってきました。Google Colabなども普及して手軽にGPUも使えるようになってきています。
一方で、PythonやGPUを用いた場合、学習は高速化できても、推論はボトルネックとなることがしばしばあります。
そこで、PythonとGPUでネットワークのパラメータを学習し、C++とCPUで推論を行うのが高速化への第一歩だと思います。
ここでは、ネットに落ちている色々な情報を統合し、2022年現在、再現性・汎用性の高い手法に絞って紹介したいと思います。
前提条件
TensorFlowを用いてPythonで学習したモデルはすでに用意できているとします。
また、モデルのパラメータやネットワークを保存した.ckptファイルはすでに作成済みであるとします。
このモデルを用いてC++で推論を実行します。
学習・推論で用いるアクセラレータ等の環境はCPU onlyまたはCPU+GPUを想定しています。
環境
Python側：tensorflow-gpu-2.4.1
C++側：tensorflow_cc-2.4.1
C++でTensorFlowの推論を実行する手順
1. 学習済みモデルの入力層・出力層の名前を取得する。
2. .ckptファイルを.pbファイルに変換する。
3. .pbファイルとtensorflow_ccを用いて、C++でTensorFlowの推論を実行する。
このページでは1.を解説します。
学習済みモデルの入力層・出力層の名前を取得する。
あとの工程で必要になるので、まずは学習済みモデルの入力層・出力層の名前を調べます。
以下のプログラムを使用します。
import tensorflow.compat.v1 as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # If you want to use only CPUs
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')
# Show Layers' Names
########################################
print("Start")
graph = sess.graph
with graph.as_default():
for op in graph.get_operations():
print(op.type, op.name, op.outputs) # type, name, shape
print("Finish")
########################################
ネットワークの構造と学習済みパラメータの読み込み
ネットワークの構造は、model.ckpt-XXXX.metaに保存されています。
また、学習済みパラメータの情報はmodel.ckpt-XXXXに保存されています。
XXXXに、は学習回数などの数値が記されています。
# Restore the graph
saver = tf.train.import_meta_graph("./checkpoints/model.ckpt-XXXX.meta")
# Load weights
saver.restore(sess, './checkpoints/model.ckpt-XXXX')
各Layerの情報の出力
get_operations()を使って、各Layerのtype （型）, name （名前）, shape （次元）を取得します。
# Show Layers' Names
########################################
print("Start")
graph = sess.graph
with graph.as_default():
for op in graph.get_operations():
print(op.type, op.name, op.outputs) # type, name, shape
print("Finish")
########################################
出力結果が以下になります。
Start
Placeholder Placeholder [<tf.Tensor 'Placeholder:0' shape=(4, 20, 16, 16, 16, 8) dtype=float32>]
Placeholder Placeholder_1 [<tf.Tensor 'Placeholder_1:0' shape=(4, 18, 16, 16, 16, 8) dtype=float32>]
Placeholder Placeholder_2 [<tf.Tensor 'Placeholder_2:0' shape=() dtype=float32>]
Const XXXXXXX/strided_slice/stack [<tf.Tensor 'XXXXXXX/strided_slice/stack:0' shape=(2,) dtype=int32>]
Const XXXXXXX/strided_slice/stack_1 [<tf.Tensor 'XXXXXXX/strided_slice/stack_1:0' shape=(2,) dtype=int32>]
Const XXXXXXX/strided_slice/stack_2 [<tf.Tensor 'XXXXXXX/strided_slice/stack_2:0' shape=(2,) dtype=int32>]
StridedSlice XXXXXXX/strided_slice [<tf.Tensor 'XXXXXXX/strided_slice:0' shape=(4, 16, 16, 16, 8) dtype=float32>]
Const XXXXXXX/zeros/shape_as_tensor [<tf.Tensor 'XXXXXXX/zeros/shape_as_tensor:0' shape=(5,) dtype=int32>]
Const XXXXXXX/zeros/Const [<tf.Tensor 'XXXXXXX/zeros/Const:0' shape=() dtype=float32>]
Fill XXXXXXX/zeros [<tf.Tensor 'XXXXXXX/zeros:0' shape=(4, 16, 16, 16, 32) dtype=float32>]
Const XXXXXXX/zeros_1/shape_as_tensor [<tf.Tensor 'XXXXXXX/zeros_1/shape_as_tensor:0' shape=(5,) dtype=int32>]
(中略)
Finish
出力結果のdtypeやshapeの値をヒントにして、入力層・出力層に相当するLayerに見当をつけます。
outputは数千~数万行に及ぶ場合もあるので、python FindLayer.py |& tee layer.logなどと実行して、ログを保存してから名前を調べると良いかもしれません。
Layerの名前を調べるのは非常に根気のいる作業です。
次の手順
2. .ckptファイルを.pbファイルに変換する。
(追記しました。)