Keras/TensorFlow で作成したモデルの学習を C++ で行う

はじめに

Python で Keras/TensorFlow を使って初期状態のモデルの作成を行い，C++ を使ってそのモデルの学習を行ってみたいと思います。

想定されるケースは，レアだとは思いますが，僅かでも早く学習を行いたい場合や，学習用のマシンに Python の環境をインストールすることが困難な場合を想定しています。

今回作成したコードは github に置いてあるので，詳細はこちらをご確認ください。

実行環境

Windows
Python 3.6
keras 2.2.4
tensorflow 1.10.0 (*)
Visual Studio 2015

(*) Python と C++ で使用する TensorFlow のバージョンは揃えていないとエラーが発生する場合があるようです。

C++ 用の tensorflow.dll はこちらのサイトからダウンロードしました。
- GitHub - fo40225/tensorflow-windows-wheel: Tensorflow prebuilt binary for Windows

参考にしたサイト

基本的な流れ

Keras + TensorFlow を用いてモデルを作成する (Python)
モデルをロードして学習を行う (C++)

Computation Graph を作成する (Python)

まず、Keras + TensorFlow で Computation Graph を作成します。
ここでは入力層の名前を "input", 出力層の名前を "output", 学習用の正解データの入力を "target" としています。この名前は後ほど C++ で学習を行う際に必要になります。

num_layer = 3
input_dim = 784
n_hidden = 200
num_classes = 10
dropout = 0.2

# モデル作成
model = Sequential()

# 入力層
model.add(InputLayer(input_shape=(input_dim,), name='input'))
model.add(Dense(n_hidden, kernel_initializer=weight_variable))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(Dropout(dropout))

# 中間層
for i in range(num_layer):
    model.add(Dense(n_hidden, kernel_initializer=weight_variable))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Dropout(dropout))

# 出力層
model.add(Dense(num_classes, kernel_initializer=weight_variable))
model.add(Activation('softmax', name='output'))

model.summary()

x = tf.placeholder(tf.float32, shape=[None, input_dim], name='image')
y_ = tf.placeholder(tf.float32, shape=[None, num_classes], name='target')
y = model(x)

loss = tf.losses.mean_squared_error(y, y_)
optimizer = tf.train.AdamOptimizer()

train_op = optimizer.minimize(loss, name='train')

init = tf.global_variables_initializer()

saver_def = tf.train.Saver().as_saver_def()

次に model.pb と model.meta というファイル名でモデルを出力します(*)。 model.pb は C++ で学習を行う際に使用します。
model.meta は Python でモデルを freeze する際に使用します。

(*) C++ で model.meta を読み込む、もしくは Python で model.pb を読み込むことができればどちらか一方の出力でよいはずですが，方法が見つからず断念しました。

# .meta の出力
saver = tf.train.Saver()
saver.export_meta_graph('model.meta')

# .pb の出力
with open('graph.pb', 'wb') as f:
    f.write(tf.get_default_graph().as_graph_def().SerializeToString())

C++ で学習を行う

まず、 model.pb の読み込みを行い、session を起動します。

const string graph_def_filename = "model.pb";

// Setup global state for TensorFlow.
tensorflow::port::InitMain(argv[0], &argc, &argv);

tensorflow::GraphDef graph_def;
TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
                                        graph_def_filename, &graph_def));
std::unique_ptr<tensorflow::Session> session(tensorflow::NewSession(tensorflow::SessionOptions()));
TF_CHECK_OK(session->Create(graph_def));

次に checkpoint 用のフォルダを指定し，フォルダ内に既存の checkpoint があればそれを読み込みます。無ければ session はクリアされます。

const string checkpoint_dir = "./checkpoints";
const string checkpoint_prefix = checkpoint_dir + "/checkpoint";

if (directory_exists(checkpoint_dir)) {
  std::cout << "Restoring model weights from checkpoint\n";
  tensorflow::Tensor t(tensorflow::DT_STRING, tensorflow::TensorShape());
  t.scalar<string>()() = checkpoint_prefix;
  TF_CHECK_OK(session->Run({{"save/Const", t}}, {}, {"save/restore_all"}, nullptr));
} else {
  std::cout << "Initializing model weights\n";
  TF_CHECK_OK(session->Run({}, {}, {"init"}, nullptr));
}

学習用の画像データと正解ラベルを読み込み，モデルの学習を回します。

auto train_x = read_training_file("MNIST_data/train-images.idx3-ubyte");
auto train_y = read_label_file("MNIST_data/train-labels.idx1-ubyte");

for (int i = 0; i < 20; ++i) {
  std::cout << "Epoch: " << i << std::endl;
  run_train_step(session, train_x, train_y);
}

TensorFlow を呼び出して学習を行う処理は以下のようになります。

"image" には (画像データ数, 784) の Tensor を， "target" には (画像データ数, 10) の Tensor を指定しています。"image" と "target" は Computation Graph を作成するときに placeholder で指定した名前となります。

void run_train_step(const std::unique_ptr<tensorflow::Session>& session, const std::vector<vector<float>>& input_batch,
                    const std::vector<float>& target_batch) {
  auto train_y = to_one_hot(target_batch);
  vector<std::pair<string, tensorflow::Tensor>> inputs = {
    {"image", MakeTensor(input_batch)},
    {"target", MakeTargetTensor(train_y)}
  };
  TF_CHECK_OK(session->Run(inputs, {}, {"train"}, nullptr));
}

std::vector から Tensor への変換は以下のように行っています。

// 画像データの変換
tensorflow::Tensor MakeTensor(const std::vector<vector<float>>& batch) {
  tensorflow::Tensor t(tensorflow::DT_FLOAT,
                       tensorflow::TensorShape({(int)batch.size(), 784}));
  auto dst = t.flat<float>().data();
  for (auto img : batch) {
    std::copy_n(img.begin(), 784, dst);
    dst += 784;
  }
  return t;
}

// 正解データの変換
tensorflow::Tensor MakeTargetTensor(const std::vector<vector<float>>& batch) {
  tensorflow::Tensor t(tensorflow::DT_FLOAT,
                       tensorflow::TensorShape({(int)batch.size(), 10}));
  auto dst = t.flat<float>().data();
  for (auto target : batch) {
    std::copy_n(target.begin(), 10, dst);
    dst += 10;
  }
  return t;
}

最後に学習結果の checkpoint を保存します。

  tensorflow::Tensor t(tensorflow::DT_STRING, tensorflow::TensorShape());
  t.scalar<string>()() = checkpoint_prefix;
  TF_CHECK_OK(session->Run({{"save/Const", t}}, {}, {"save/control_dependency"}, nullptr));

今回はこれで終了です。
github のコードでは学習の前後で推論を行っているので，実行すると checkpoints ディレクトリが無い初期状態だと精度が 10% 前後，学習後だと精度が 90% 前後になることが確認できると思います。

モデルの freeze と freeze されたモデルを使用した推論についてはまた次回。