このページは Cloud Translation API によって翻訳されました。

AOT コンパイルの使用

tfcompile とは

tfcompile は、事前（AOT）で TensorFlow グラフを実行可能コードにコンパイルするスタンドアロンツールです。これにより、バイナリの合計サイズを削減でき、ランタイムオーバーヘッドも回避できます。tfcompile の一般的なユースケースは、推論グラフをモバイルデバイス用の実行コードにコンパイルすることです。

TensorFlow グラフは通常、TensorFlow ランタイムによって実行されます。これにより、グラフ内の各ノードの実行でランタイムオーバーヘッドが発生します。グラフ自体に加えて TensorFlow ランタイムのコードも利用できる必要があるため、合計バイナリサイズも大きくなります。tfcompile によって生成された実行可能コードは TensorFlow ランタイムを使用せず、実際に計算で使用されるカーネルにのみ依存関係があります。

コンパイラは XLA フレームワーク上に構築されています。TensorFlow を XLA フレームワークにブリッジするコードは、tensorflow/compiler にあります。

tfcompile の内容

tfcompile は、TensorFlow のフィードとフェッチのコンセプトで識別されるサブグラフを受け取り、そのサブグラフを実装する関数を生成します。feeds は関数の入力引数で、fetches は関数の出力引数です。すべての入力はフィードで完全に指定する必要があります。結果として得られるプルーニングされたサブグラフに、プレースホルダノードや変数ノードを含めることはできません。すべてのプレースホルダと変数をフィードとして指定するのが一般的です。これにより、生成されるサブグラフにこれらのノードが含まれなくなります。生成された関数は、関数のシグネチャをエクスポートするヘッダーファイルと、実装を含むオブジェクトファイルとともに、cc_library としてパッケージ化されます。ユーザーは、生成された関数を適宜呼び出すコードを記述します。

tfcompile の使用

このセクションでは、TensorFlow サブグラフから tfcompile を使用して実行可能バイナリを生成する手順の概要を説明します。ステップは次のとおりです。

ステップ 1: コンパイルするサブグラフを構成する
ステップ 2: tf_library ビルドマクロを使用してサブグラフをコンパイルする
ステップ 3: サブグラフを呼び出すコードを記述する
ステップ 4: 最終バイナリを作成する

ステップ 1: コンパイルするサブグラフを構成する

生成された関数の入力引数と出力引数に対応するフィードとフェッチを特定します。次に、tensorflow.tf2xla.Config プロトコルで feeds と fetches を構成します。

# Each feed is a positional input argument for the generated function.  The order
# of each entry matches the order of each input argument.  Here “x_hold” and “y_hold”
# refer to the names of placeholder nodes defined in the graph.
feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}

# Each fetch is a positional output argument for the generated function.  The order
# of each entry matches the order of each output argument.  Here “x_y_prod”
# refers to the name of a matmul node defined in the graph.
fetch {
  id { node_name: "x_y_prod" }
}

ステップ 2: tf_library ビルドマクロを使用してサブグラフをコンパイルする

このステップでは、tf_library ビルドマクロを使用して、グラフを cc_library に変換します。cc_library は、グラフから生成されたコードを含むオブジェクトファイルと、生成されたコードにアクセスするためのヘッダーファイルで構成されます。tf_library は、tfcompile を利用して TensorFlow グラフを実行可能コードにコンパイルします。

load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# Use the tf_library macro to compile your graph into executable code.
tf_library(
    # name is used to generate the following underlying build rules:
    # <name>           : cc_library packaging the generated header and object files
    # <name>_test      : cc_test containing a simple test and benchmark
    # <name>_benchmark : cc_binary containing a stand-alone benchmark with minimal deps;
    #                    can be run on a mobile device
    name = "test_graph_tfmatmul",
    # cpp_class specifies the name of the generated C++ class, with namespaces allowed.
    # The class will be generated in the given namespace(s), or if no namespaces are
    # given, within the global namespace.
    cpp_class = "foo::bar::MatMulComp",
    # graph is the input GraphDef proto, by default expected in binary format.  To
    # use the text format instead, just use the ‘.pbtxt’ suffix.  A subgraph will be
    # created from this input graph, with feeds as inputs and fetches as outputs.
    # No Placeholder or Variable ops may exist in this subgraph.
    graph = "test_graph_tfmatmul.pb",
    # config is the input Config proto, by default expected in binary format.  To
    # use the text format instead, use the ‘.pbtxt’ suffix.  This is where the
    # feeds and fetches were specified above, in the previous step.
    config = "test_graph_tfmatmul.config.pbtxt",
)

この例の GraphDef プロトコル（test_graph_tfmatmul.pb）を生成するには、make_test_graphs.py を実行し、--out_dir フラグを使用して出力場所を指定します。

一般的なグラフには、トレーニングによって学習された重みを表す Variables が含まれますが、tfcompile は Variables を含むサブグラフをコンパイルできません。freeze_graph.py ツールは、チェックポイントファイルに保存されている値を使用して、変数を定数に変換します。便宜上、tf_library マクロは freeze_checkpoint 引数をサポートしています。これにより、ツールが実行されます。その他の例については、tensorflow/compiler/aot/tests/BUILD をご覧ください。

コンパイルされたサブグラフに表示される定数は、生成されたコードに直接コンパイルされます。生成された関数に定数を渡すには、定数をコンパイルする代わりにフィードとして渡します。

tf_library ビルドマクロの詳細については、tfcompile.bzl をご覧ください。

基盤となる tfcompile ツールの詳細については、tfcompile_main.cc をご覧ください。

ステップ 3: サブグラフを呼び出すコードを記述する

このステップでは、前のステップで tf_library ビルドマクロによって生成されたヘッダーファイル（test_graph_tfmatmul.h）を使用して、生成されたコードを呼び出します。ヘッダーファイルは、ビルドパッケージに対応する bazel-bin ディレクトリにあり、tf_library ビルドマクロに設定された name 属性に基づいて名前が付けられます。たとえば、test_graph_tfmatmul に対して生成されるヘッダーは test_graph_tfmatmul.h になります。以下は、生成される内容の簡略版です。生成されたファイル（bazel-bin）には、他にも有用なコメントが含まれています。

namespace foo {
namespace bar {

// MatMulComp represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code.
class MatMulComp {
 public:
  // AllocMode controls the buffer allocation mode.
  enum class AllocMode {
    ARGS_RESULTS_AND_TEMPS,  // Allocate arg, result and temp buffers
    RESULTS_AND_TEMPS_ONLY,  // Only allocate result and temp buffers
  };

  MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS);
  ~MatMulComp();

  // Runs the computation, with inputs read from arg buffers, and outputs
  // written to result buffers. Returns true on success and false on failure.
  bool Run();

  // Arg methods for managing input buffers. Buffers are in row-major order.
  // There is a set of methods for each positional argument.
  void** args();

  void set_arg0_data(float* data);
  float* arg0_data();
  float& arg0(size_t dim0, size_t dim1);

  void set_arg1_data(float* data);
  float* arg1_data();
  float& arg1(size_t dim0, size_t dim1);

  // Result methods for managing output buffers. Buffers are in row-major order.
  // Must only be called after a successful Run call. There is a set of methods
  // for each positional result.
  void** results();


  float* result0_data();
  float& result0(size_t dim0, size_t dim1);
};

}  // end namespace bar
}  // end namespace foo

生成された C++ クラスは、tf_library マクロで指定された cpp_class であったため、foo::bar 名前空間で MatMulComp と呼ばれます。生成されるすべてのクラスは同様の API を持ちますが、引数と結果バッファを処理するメソッドだけが異なります。これらのメソッドは、tf_library マクロの feed 引数と fetch 引数で指定されたバッファの数と型によって異なります。

生成されたクラス内で管理されるバッファは 3 種類あります。args は入力を表し、results は出力を表します。temps は計算を実行するために内部で使用される一時バッファを表します。デフォルトでは、生成されたクラスの各インスタンスがこれらのバッファをすべて割り当て、管理します。AllocMode コンストラクタ引数を使用すると、この動作を変更できます。すべてのバッファは 64 バイト境界に揃えられます。

生成される C++ クラスは、XLA によって生成された低レベルコードのラッパーにすぎません。

tfcompile_test.cc に基づいて生成された関数を呼び出す例:

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "third_party/tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());


  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

ステップ 4: 最終バイナリを作成する

このステップでは、ステップ 2 で tf_library が生成したライブラリと、ステップ 3 で記述したコードを組み合わせ、最終的なバイナリを作成します。bazel BUILD ファイルの例を次に示します。

# Example of linking your binary
# Also see //tensorflow/compiler/aot/tests/BUILD
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# The same tf_library call from step 2 above.
tf_library(
    name = "test_graph_tfmatmul",
    ...
)

# The executable code generated by tf_library can then be linked into your code.
cc_binary(
    name = "my_binary",
    srcs = [
        "my_code.cc",  # include test_graph_tfmatmul.h to access the generated header
    ],
    deps = [
        ":test_graph_tfmatmul",  # link in the generated object file
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)