Dataset を使った TensorFlow のモデルを SavedModel 形式のファイルから復元して推論を行う（C++編）

はじめに

前回の記事 では Python で TensorFlow の Dataset を使用して学習を行い、モデルを SavedModel の形式で保存しました。

今回は、そのモデルを C++ のコードから読み込んで推論を実行してみたいと思います。

本記事で説明しているソースコードは こちら に載せてあります。

実行環境

• Windows
• Visual Studio 2015
• tensorflow 1.10.0

モデルの読み込み

初めにモデルを読み込みます。

infer.cpp

  const std::string export_dir = "./model";
  SavedModelBundle bundle;
  SessionOptions session_options;
  RunOptions run_options;

  Status status = tensorflow::LoadSavedModel(session_options, run_options, export_dir, {tensorflow::kSavedModelTagServe}, &bundle);
  if (!status.ok()) {
    std::cout << status.ToString() << std::endl;
    return -1;
  }

入力データの作成

入力用データ（Tensor 型）を作成します。

infer.cpp

  Tensor batch_size(tensorflow::DT_INT64, tensorflow::TensorShape());
  auto dst = batch_size.flat<long long>().data();
  long long bsize = 3L;
  memcpy(dst, &bsize, sizeof(bsize));

  Tensor input(tensorflow::DT_FLOAT, TensorShape({3, 1}));
  auto input_dst = input.flat<float>().data();
  float arr[3] = {2.0, 3.0, 4.0};
  memcpy(input_dst, arr, sizeof(arr));

Iterator の初期化

Iterator の初期化を行います。"dataset_init" はモデル作成時に指定した initializer の名前です。
また、"input", "target", "batch_size" は同様にモデル作成時に指定した placeholder の名前です。

infer.cpp

  status = bundle.session->Run(
      {{"input", input}, {"target", input}, {"batch_size", batch_size}},
      {},
      {"dataset_init"},
      nullptr);
  if (!status.ok()) {
    std::cout << "Failed to run sesssion (dataset_init)" << std::endl;
    std::cout << status.ToString() << std::endl;
    return -1;
  }

推論の実行

最後に推論を実行します。"output:0" はモデル作成時に指定した出力ノードの名前です。

infer.cpp

  std::vector<Tensor> outputs;
  status = bundle.session->Run({}, {"output:0"}, {}, &outputs);
  if (!status.ok()) {
    std::cout << "Failed to run sesssion (output:0)" << std::endl;
    std::cout << status.ToString() << std::endl;
    return -1;
  }

  Tensor a = outputs.at(0);
  const int out_dim = 3;
  for (int i = 0; i < out_dim; i++) {
    std::cout << a.flat<float>()(i) << std::endl;
  }

実行すると以下のように Python 版とほぼ同じ推論結果が得られます。

2 4.16
3 6.12808
4 8.09616