この記事について
Tensorflow XLAのツールであるtfcompileにより、KerasモデルのAOTコンパイルしてライブラリを生成する方法と、生成されたライブラリをアプリケーションで利用する方法について記載します。
この記事の目標
- tfcompileの実行環境を構築します
- KerasモデルをtfcompileによりAOTコンパイルします
- コンパイルしたライブラリを利用するC++プログラムを作成します
- 作成したC++プログラムをビルドして実行します(x86マシン)
ubuntu 16.04の環境で進めていきます。
準備
tfcompileの動作のため、以下のツールを導入します。
Bazelのインストール
JDK 8をインストールします。
$sudo apt-get install openjdk-8-jdk
リポジトリを追加します。
$echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
$curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
aptで取得します。
$sudo apt-get update
$sudo apt-get install bazel
bazelの利用準備が整いました。
$bazel version
Build label: 0.13.1
Protocol Buffersのインストール
tfcompileの入力ファイルの構築にprotoc(>ver3.0)を利用します。apt-getで取得できるものはバージョンが古い場合があるため、リポジトリより直接取得します。
$wget https://github.com/google/protobuf/releases/download/v3.2.0/protoc-3.2.0-linux-x86_64.zip
解凍します。
$unzip protoc-3.2.0-linux-x86_64.zip -d protoc
パスの通っている場所に実行モジュールとインクルードヘッダをコピーします。
$sudo mv protoc/bin/* /usr/bin
$sudo mv protoc/include/* /usr/include
これでprotocのインストールは完了です。
$protoc --version
libprotoc 3.2.0
tfcompileのインストール
tfcompileのビルド
tensorflowのリポジトリよりソースコードを入手します。
$cd /path/to/workspace
$git clone --depth=1 --single-branch https://github.com/tensorflow/tensorflow
コンフィグを実行します。
$cd /path/to/workspace/tensorflow/tensorflow
$./configure
必要に応じて、Yes/Noを設定してください。最低限、下記はYesで設定します。
Do you wish to build TensorFlow with XLA JIT support? [y/N]: No XLA JIT support will be enabled for TensorFlow.
tfcompileをビルドします。(CPUモードです)
$bazel build --config=opt --verbose_failures //tensorflow/compiler/aot:tfcompile
成功するとtfcompileを使えるようになります。
$cd /path/to/tensorflow
$./bazel-bin/tensorflow/compiler/aot/tfcompile --help
tf2xla_pb2.pyの生成
protocでtf2xla.protoファイル(protocl buffers定義ファイル)をコンパイルし、pythonコードを生成します。
$cd /path/to/tensorflow/tensorflow
$protoc compiler/tf2xla/tf2xla.proto --python_out=.
tf2xla.protoと同じ場所にtf2xla_pb2.pyが生成されます。
$ls /path/to/tensorflow/tensorflow/compiler/tf2xla
tf2xla_pb2.py
このpythonモジュールは、tfcompileで利用するconfigファイルの作成補助に利用します。
モデルのコンパイル
公式ドキュメントの流れに沿って作成していきます。
Step 1/4 : graph, configに指定するファイルの作成
tfcompileはモデル実体であるサブグラフ(pb)と、モデル入出力形式を指定するためのconfigファイル(pbtxt)を必要とします。これらの情報により、tfcompileにモデルがどのような形(shape)をしているのかを正確に指定することができ、最適化が行えるようになっています。
ここでは、サブグラフ、configファイルをKeras付属のモデルを使って生成していきます。そのためのpythonプログラムを作成します。
作業用のワークフォルダを作成し、tf2xla_pb2.pyをコピーしてきます。
$cd /path/to/tensorflow
$mkdir myaot
$cp tensorflow/compiler/tf2xla/tf2xla_pb2.py myaot/
create_graph.py
#-------------------------------------------------
# Kerasサンプルのモデルを取得します
#-------------------------------------------------
import tensorflow as tf
model = tf.keras.applications.ResNet50()
#-------------------------------------------------
# グラフを保存します (test_graph.pb)
#-------------------------------------------------
ss = tf.keras.backend.get_session()
output_node_names = [node.op.name for node in model.outputs]
graphdef = tf.graph_util.convert_variables_to_constants(ss, ss.graph_def, output_node_names)
tf.train.write_graph(graphdef, '.', 'test_graph.pb', as_text=False)
#-------------------------------------------------
# configファイルを作成します
# (test_graph.config.pbtxt)
#-------------------------------------------------
import tf2xla_pb2
# configの取得
config = tf2xla_pb2.Config()
# feed (入力形式)の指定
batch_size = 1
for x in model.inputs:
# shapeを[1,Dim(224),Dim(224),Dim(3)] (ResNet50の入力形式)にセット
x.set_shape([batch_size] + list(x.shape)[1:])
feed = config.feed.add()
feed.id.node_name = x.op.name
feed.shape.MergeFrom(x.shape.as_proto())
# fetch (出力形式)の指定
for x in model.outputs:
fetch = config.fetch.add()
fetch.id.node_name = x.op.name
# configファイルの保存
with open('test_graph.config.pbtxt', 'w') as fo:
out_txt = str(config)
# print(out_txt) # for display
fo.write(out_txt)
以下の二つのファイルが生成されます。
$python create_graph.py
test_graph.pb
test_graph.config.txt
グラフの構成をコンパイラに指定するためのconfigファイルは以下のような内容となります。
test_graph.config.pbtxt
feed {
id {
node_name: "input_1"
}
shape {
dim {
size: 1
}
dim {
size: 224
}
dim {
size: 224
}
dim {
size: 3
}
}
}
fetch {
id {
node_name: "fc1000/Softmax"
}
}
modelの形式との対応は以下の通りです。
model.summary()
________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
================================================================================
input_1 (InputLayer) (None, 224, 224, 0
________________________________________________________________________________
・
・
・
________________________________________________________________________________
fc1000 (Dense) (None, 1000) 2049000 flatten_1[0][0]
================================================================================
Step 2/4 : サブグラフのコンパイル
step1で作成したサブグラフをtfcompileによりコンパイルします。ここでは公式ドキュメントに沿って、tfcompileを直接利用するのではなく、Bazelを通して実行します。
tfcompileを利用するためのBUILDファイルを作成します。
$vi myaot/BUILD
BUILDファイルを以下のように修正します。ここで定義した内容は最終的にtfcompile実行時の引数に渡されます。
myaot/BUILD
tf_library(
name = 'test_graph',
config = 'test_graph.config.pbtxt',
cpp_class = 'TestGraph',
graph = 'test_graph.pb',
)
bazel buildを実行します。
$cd /path/to/tensorflow
$bazel build //myaot:test_graph
成功すると、以下のようなメッセージが表示され、
INFO: Build completed successfully, 2671 total actions
ライブラリとヘッダファイルが生成されます。
$ls bazel-genfiles/myaot/
test_graph.h test_graph_tfcompile_function.o test_graph_tfcompile_metadata.o
生成されたインクルードヘッダは以下のようになります。このファイルは自動生成されるため、編集することはありません。
test_graph.h
// Generated by tfcompile, the TensorFlow graph compiler. DO NOT EDIT!
//
// This header was generated via ahead-of-time compilation of a TensorFlow
// graph. An object file corresponding to this header was also generated.
// This header gives access to the functionality in that object file.
//
// clang-format off
#ifndef TFCOMPILE_GENERATED___myaot__test_graph_H_ // NOLINT(build/header_guard)
#define TFCOMPILE_GENERATED___myaot__test_graph_H_ // NOLINT(build/header_guard)
#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
#include "tensorflow/core/platform/types.h"
namespace Eigen { struct ThreadPoolDevice; }
namespace xla { class ExecutableRunOptions; }
// (Implementation detail) Entry point to the function in the object file.
extern "C" void __myaot__test_graph(
void* result, const xla::ExecutableRunOptions* run_options,
const void** args, void** temps, tensorflow::int64* profile_counters);
// TestGraph represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code. This extends the generic
// XlaCompiledCpuFunction class with statically type-safe arg and result
// methods. Usage example:
//
// TestGraph computation;
// // ...set args using computation.argN methods
// CHECK(computation.Run());
// // ...inspect results using computation.resultN methods
//
// The Run method invokes the actual computation, with inputs read from arg
// buffers, and outputs written to result buffers. Each Run call may also use
// a set of temporary buffers for the computation.
//
// By default each instance of this class manages its own arg, result and temp
// buffers. The AllocMode constructor parameter may be used to modify the
// buffer allocation strategy.
//
// Under the default allocation strategy, this class is thread-compatible:
// o Calls to non-const methods require exclusive access to the object.
// o Concurrent calls to const methods are OK, if those calls are made while it
// is guaranteed that no thread may call a non-const method.
//
// The logical function signature is:
// (arg0: f32[1,224,224,3]) -> (f32[1,1000])
//
// Memory stats:
// arg bytes total: 602112
// arg bytes aligned: 602112
// temp bytes total: 17815208
// temp bytes aligned: 17815296
class TestGraph : public tensorflow::XlaCompiledCpuFunction {
public:
// Number of input arguments for the compiled computation.
static constexpr size_t kNumArgs = 1;
// Byte size of each argument buffer. There are kNumArgs entries.
static const intptr_t* ArgSizes() {
static constexpr intptr_t kArgSizes[kNumArgs] = {602112};
return kArgSizes;
}
// Returns static data used to create an XlaCompiledCpuFunction.
static const tensorflow::XlaCompiledCpuFunction::StaticData& StaticData() {
static XlaCompiledCpuFunction::StaticData* kStaticData = [](){
XlaCompiledCpuFunction::StaticData* data =
new XlaCompiledCpuFunction::StaticData;
data->raw_function = __myaot__test_graph;
data->arg_sizes = ArgSizes();
data->num_args = kNumArgs;
data->temp_sizes = TempSizes();
data->num_temps = kNumTemps;
data->result_index = kResultIndex;
data->arg_names = StaticArgNames();
data->result_names = StaticResultNames();
data->program_shape = StaticProgramShape();
data->hlo_profile_printer_data = StaticHloProfilePrinterData();
return data;
}();
return *kStaticData;
}
TestGraph(AllocMode alloc_mode = AllocMode::ARGS_RESULTS_PROFILES_AND_TEMPS)
: XlaCompiledCpuFunction(StaticData(), alloc_mode) {}
TestGraph(const TestGraph&) = delete;
TestGraph& operator=(const TestGraph&) = delete;
// Arg methods for managing input buffers. Buffers are in row-major order.
// There is a set of methods for each positional argument, with the following
// general form:
//
// void set_argN_data(void* data)
// Sets the buffer of type T for positional argument N. May be called in
// any AllocMode. Must be called before Run to have an affect. Must be
// called in AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY for each positional
// argument, to set the argument buffers.
//
// T* argN_data()
// Returns the buffer of type T for positional argument N.
//
// T& argN(...dim indices...)
// Returns a reference to the value of type T for positional argument N,
// with dim indices specifying which value. No bounds checking is performed
// on dim indices.
void set_arg0_data(void* data) {
set_arg_data(0, data);
}
float* arg0_data() {
return static_cast<float*>(arg_data(0));
}
float& arg0(size_t dim0, size_t dim1, size_t dim2, size_t dim3) {
return (*static_cast<float(*)[1][224][224][3]>(
arg_data(0)))[dim0][dim1][dim2][dim3];
}
const float* arg0_data() const {
return static_cast<const float*>(arg_data(0));
}
const float& arg0(size_t dim0, size_t dim1, size_t dim2, size_t dim3) const {
return (*static_cast<const float(*)[1][224][224][3]>(
arg_data(0)))[dim0][dim1][dim2][dim3];
}
// 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, with the following general form:
//
// T* resultN_data()
// Returns the buffer of type T for positional result N.
//
// T& resultN(...dim indices...)
// Returns a reference to the value of type T for positional result N,
// with dim indices specifying which value. No bounds checking is performed
// on dim indices.
//
// Unlike the arg methods, there is no set_resultN_data method. The result
// buffers are managed internally, and may change after each call to Run.
float* result0_data() {
return static_cast<float*>(result_data(0));
}
float& result0(size_t dim0, size_t dim1) {
return (*static_cast<float(*)[1][1000]>(
result_data(0)))[dim0][dim1];
}
const float* result0_data() const {
return static_cast<const float*>(result_data(0));
}
const float& result0(size_t dim0, size_t dim1) const {
return (*static_cast<const float(*)[1][1000]>(
result_data(0)))[dim0][dim1];
}
private:
// Number of result and temporary buffers for the compiled computation.
static constexpr size_t kNumTemps = 10;
// The 0-based index of the result tuple in the temporary buffers.
static constexpr size_t kResultIndex = 2;
// Byte size of each result / temporary buffer. There are kNumTemps entries.
static const intptr_t* TempSizes() {
static constexpr intptr_t kTempSizes[kNumTemps] = {-1, 4000, 8, -1, -1, -1, -1, -1, -1, 17811200};
return kTempSizes;
}
// Array of names of each positional argument, terminated by nullptr.
static const char** StaticArgNames() {
return nullptr;
}
// Array of names of each positional result, terminated by nullptr.
static const char** StaticResultNames() {
return nullptr;
}
// Shape of the args and results.
static const xla::ProgramShape* StaticProgramShape() {
static const xla::ProgramShape* kShape = nullptr;
return kShape;
}
// Metadata that can be used to pretty-print profile counters.
static const xla::HloProfilePrinterData* StaticHloProfilePrinterData() {
static const xla::HloProfilePrinterData* kHloProfilePrinterData =
nullptr;
return kHloProfilePrinterData;
}
};
#endif // TFCOMPILE_GENERATED___myaot__test_graph_H_
// clang-format on
Step 3/4 : アプリケーションコードの作成
生成されたクラスを利用するコードを作成します。生成されたクラスの使い方についてはヘッダファイル内に説明が記載されています。
my_code.cpp
#include <iostream>
#include <fstream>
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL
#include "myaot/test_graph.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
// ビットマップ画像をロードするクラス
class RawImage
{
public:
RawImage(){}
bool loadImage(std::string filename){
// ファイル読み込み
std::ifstream fin( filename.c_str(), std::ios::in | std::ios::binary );
if (!fin){
std::cout << "cannot open file:" << filename << std::endl;
return false;
}
// ビットマップファイルのヘッダを読み飛ばします(54byte)
for(int i=0; i<54; i++){
char dmy;
fin.read(&dmy, sizeof(dmy) );
}
// 画像データを1byteずつ読み込みます
// ResNet50モデルの入力に使えるように前処理も行います
// ビットマップ形式ではY軸は上下反転、カラー情報はBGR形式で格納されているため
// 画像情報を左下から読み込んで、順番に格納していきます
for(int y=223; y>=0; y--){
for(int x=0; x<224; x++ ){
for(int c=0; c<3; c++){
if(!fin.eof()){
// 1byteずつB->G->Rの順番で読み込みます
unsigned char v;
fin.read( (char*)&v, sizeof(v) );
// 前処理として画像の画素値よりVGG16の平均画素値を減算します
// (学習時と同じ前処理を実施)
float vv = (float)v;
const float mean[3] = {103.939,116.779,123.68}; // BGRのそれぞれの平均値
vv = vv - mean[c];
data[0][y][x][c] = vv;
}
}
}
}
fin.close();
}
// 入力画像の前処理済みのデータ
float data[1][224][224][3];
};
// 推論実行関数
int run(const float *p_input, int input_size, float *p_output, int output_size) {
Eigen::ThreadPool tp(std::thread::hardware_concurrency());
Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
Graph computation;
computation.set_thread_pool(&device);
// 入力値のセット
computation.set_arg0_data((void*)p_input);
// 推論実行
auto success = computation.Run();
if(not success){
return -1;
}
// 推論結果を取得
std::copy(computation.result0_data(), computation.result0_data() + output_size/sizeof(float), p_output);
return 0;
}
// 配列内の最大値となる添え字を返す
int max( float* array, int len ){
int max_i = 0;
for(int i=0; i<len; i++){
if( array[i] > array[max_i] ){
max_i = i;
}
}
return max_i;
}
// Main関数
int main(){
RawImage rawimg; // ビットマップ画像
float result[1][1000]; // 結果取得用
// 24bit color 224x224のビットマップイメージを取得
rawimg.loadImage("./cat224.bmp");
// ResNet50の推論モデルに取得したビットマップ画像を入力し結果を取得する
int ret = run( (const float*)rawimg.data, sizeof(rawimg.data), (float*) result, sizeof(result) );
// 最大のスコアを持つIDを検索する(ResNet50の分類クラスのどれにヒットしたか確認)
int max_i = max( &(result[0][0]), 1000 );
// 結果出力
std::cout << "result max_i is " << max_i << std::endl;
return 0;
}
Step 4/4 : 実行モジュールの生成
gcc/g++を直接使ってコンパイルすることもできますが、ここでは、Bazelを使った方法で実行モジュールを生成します。先ほど作成したBUILDファイルにC++ビルド用のcc_binaryルールを追記します。
myaot/BUILD
cc_binary(
name = "my_code",
srcs = [
"my_code.cc"
],
deps = [
":test_graph",
"//third_party/eigen3"
],
linkopts = [
"-lpthread"
],
)
ビルドを実行します。
$cd /path/to/tensorflow
$bazel build --verbose_failures ://myaot:my_code
成功すると、実行バイナリが生成されます。
$ls ./bazel-bin/myaot
my_code
参考ドキュメント
-
環境準備
-
Tensorflow XLA/AOT