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.を解説します。

学習済みモデルの入力層・出力層の名前を取得する。

あとの工程で必要になるので、まずは学習済みモデルの入力層・出力層の名前を調べます。
以下のプログラムを使用します。

FindLayers.py

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")
    ########################################

出力結果が以下になります。

output

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ファイルに変換する。
(追記しました。)

参考にしたサイト

Pythonで学習したネットワークをC++で実行する

ニューラルネットワークの情報のロード