Opを実装してみる
今までMatMul、Sub、Div、…などいろいろなOpがあったと思いますが、ここではOp新たに定義して使ってみたいと思います。
tensorflow/core/user_ops/fact.cc
tenserflow/user_ops/ackermann_op.cc
というexampleがありましたが、これについてはAdding New Opを見た方が良いと思うので、これを元に進めたいと思います。
1回ですべてまとめるとかなり長くなりそうだったので、多分3回くらいに分割してやっていきたいと思います。
1回目: Op実装の仕方
2回目: Attrの追加
3回目: Opのポリモーフィズム
もしかしたら２回かも
説明がわかりにくかったらごめんなさい！
1. 追加するOpのインターフェースの定義
まず初めにやることは、OpDefというものを定義して、OpRegistryというものに登録することになります。これは、以下に示したREGISTER_OPというマクロを呼んで定義することになります。
REGISTER_OP("my_op_name")
.Attr("<name>:<type>")
.Attr("<name>:<type>=<default>")
.Input("<name>:<type-expr>")
.Input("<name>:Ref(<type-expr>)")
.Output("<name>:<type-expr>")
.Doc(R"(
<1-line summary>
<rest of the description (potentially many lines)>
<name-of-attr-input-or-output>: <description of name>
<name-of-attr-input-or-output>: <description of name;
if long, indent the description on subsequent lines> )");
今回実装したOpの例を示します。
Opの名前のルールは、
1. CamelCase
2. アンダースコアから始めない
です。
今回の例は、
REGISTER_OP("My") //Opの名前は"My"
.Input("msg: string") //1つ目の入力はstringのTensorで名前は"msg"
.Input("a: int32) //2つ目の入力はint32のTensorで名前は"a"
.Output("b: int32") //出力のTensorはint32で名前は"b"
.Doc(R"doc(
Test Op.
)doc");
Input/Outputはそれぞれ複数指定可能です。
Attrについては次回詳しく。
2. Kernel部分の実装
Opの実際の処理の実装部分をカーネルといい、入力のタイプや、CPU/GPUなどを区別して複数実装できるらしいです(今回は1つ)。
やることは、単純で
1. OpKernelクラスを継承したクラスを定義する
2. void Compute(OpKernelContext* context)をオーバーライドする
、となります。
このComputeに記述する処理が、このOpの処理になります。
また、OpKernelContextとは必要なものをまとめたようなオブジェクトで、入力されたTensorや出力するTensorなどを含んでいます。
今回実装したカーネルは、後述します。
Input Tensorを取り出す
入力TensorはOpKernelContextから以下のように取り出すことができます。
// i番目ののinput tensorを取り出す
const Tensor& input_tensor = context->input(i);
この順番はREGISTER_OPで定義した順番に対応します。
Output Tensorを作る
Tensor* output_tensor = NULL;
// contextにoutput_tensorを確保する
OP_REQUIRES_OK(context,
context->allocate_output(0, input_tensor.shape(), &output_tensor));
これでcontextにoutput_tensorが確保されたので、以降output_tensorを操作することで、出力を生成することができます。
3. Kernelの登録
これもマクロを呼ぶだけです。
これは、
tnesorflow/core/framework/op_kernel.h
に定義されていて、
1. Name: カーネルの名前
2. Device: 動作するデバイス
3. TypeConstraint: タイプの制限
4. HostMemory: ホストメモリ上に存在するデータの指定
5. インスタンス化するクラス(MyOp<float>, MyOp<int32>みたいな)
などいくつか設定できるみたいです。
例を見た方が早いと思うので、以下にいくつか載せておきます。
基本は以下の形です。
REGISTER_KERNEL_BUILDER(Name("My").Device(DEVICE_CPU), MyOp);
次は、Tというタイプがfloatの場合のみ動作するカーネルの場合は、
REGISTER_KERNEL_BUILDER(Name("My")
.Device(DEVICE_CPU)
.TypeConstraint<float>("T"),
MyOp<float>);
次は、動作するデバイスはGPU、入力TensorのうちshapeというTensorがホストメモリにあるという制限をつけたカーネルを登録する場合の例です。
REGISTER_KERNEL_BUILDER(Name("My")
.Device(DEVICE_GPU)
.HostMemory("shape"),
MyOp);
4. 追加するOpの共有ライブラリを作成する
追加するOpを実装したソースは以下の場所に置いておきます。
tensorflow/core/user_ops/my_op.cc
TensorFlowをバイナリからインストールした人
ここではTensorFlowのPythonライブラリから以下の関数を使用します。
1. get_include
2. get_lib
名前からわかる通り、それぞれ、ヘッダーとライブラリへのパスを返してくれます。
$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'(略)/site-packages/tensorflow/include'
といった具合。
これを使用して、共有ライブラリを作成する手順は以下の通りです。
$ TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
$ TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())')
$ g++ -std=c++11 -shared my_op.cc -o my_op.so \
-I $TF_INC -l tensorflow_framework -L $TF_LIB \
-fPIC -Wl,-rpath $TF_LIB
TensorFlowをソースからインストールした人
BUILDファイルを置いておくことで、bazelがコンパイルしてくれます。
cc_binary(
name = "my_op.so",
srcs = ["my_op.cc"],
linkopts = [
"-Wl,-Bsymbolic",
"-lm",
],
linkshared = 1,
linkstatic = 1,
deps = [
"///tensorflow/core:framework",
],
)
$ bazel build tensorflow/core/user_ops/...
使ってみる
pythonの例はあるのですが、C++での例はないようです。
pythonの場合、共有ライブラリをロードするAPIを使用しているようです。
C++ではどのようにしたらよいだろう…
適切な方法でやれば多分bazelがヘッダファイルをuser_ops.hとして作成してくれるのですが、bazelがわからない…
作成した共有ライブラリをロードする
いろいろ探していて、以下のようなAPIを見つけました。
TF_Library* TF_LoadLibrary(const char* library_filename, TF_Status* status) {
TF_Library* lib_handle = new TF_Library;
status->status = tensorflow::LoadLibrary(
library_filename, &lib_handle->lib_handle, &lib_handle->op_list.data,
&lib_handle->op_list.length);
if (!status->status.ok()) {
delete lib_handle;
return nullptr;
}
return lib_handle;
}
(ここには他にもいろいろなAPIが実装されていて、今後使用されるようになるのかも？)
また、内部で呼ばれているtensorflow::LoadLibarayは、
Status LoadLibrary(const char* library_filename, void** result,
const void** buf, size_t* len) {
...
TF_RETURN_IF_ERROR(env->LoadLibrary(library_filename, &lib));
...
}
さらにここで呼ばれているLoadLibraryは、
Status LoadLibrary(const char* library_filename, void** handle) {
*handle = dlopen(library_filename, RTLD_NOW | RTLD_LOCAL);
if (!*handle) {
return errors::NotFound(dlerror());
}
return Status::OK();
}
本質的にはdlopenで作成したライブラリをロードして使うが、結局よくわからず保留。
既存のOpと一緒にしてしまう
bazelのことをよく知っている方なら、うまい方法があるのかもしれないが、自分は全然わからなかったので、とりあえず、既存のREGISTER_OP群の末尾に自分のを置くことにしました。
/* 他のREGISTER_OP */
// "My"というOpを登録する
REGISTER_OP("My")
.Input("msg: string")
.Input("a: int32")
.Output("b: int32")
.Doc(R"doc(
MyOp
)doc");
これによって、bazelがmath_ops.hに自分のOpも含めてくれました。
本当は、user_ops.hに記述されるようにBUILDを記述するのが適切だと思います。
今回実装したカーネルを以下に載せておきます。
class MyOp : public OpKernel {
public:
explicit MyOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override {
// 1番目の入力を取得する
const Tensor& msg_tensor = context->input(0);
auto msg = msg_tensor.flat<string>();
std::cout << msg(0) << std::endl;
// 2番目の入力を取得する
const Tensor& input_tensor = context->input(1);
auto input = input_tensor.flat<int32>();
// 出力のTensorを確保する
Tensor* output_tensor = NULL;
OP_REQUIRES_OK(context,
context->allocate_output(0, input_tensor.shape(), &output_tensor));
auto output = output_tensor->template flat<int32>();
// 入力の要素を2倍する
for (int i = 0; i < input.size(); ++i)
output(i) = input(i)*2;
}
};
// "My"というカーネルを作る
REGISTER_KERNEL_BUILDER(Name("My").Device(DEVICE_CPU), MyOp);
次は、これを呼び出す側の処理について説明します。
まず、Computation Graphは以下の通りです。
Node* msg = Const((string)"my op!", b.opts()); // stringのTensor
Node* a = Const({3, 2}, b.opts()); // int32のTensor
// 定義通り、1つ目の入力にstringの、2つ目の入力にint32のTensorを渡す
My(msg, a, b.opts().WithName("my_ops_out")); // ノード名はmy_ops_out
Sessionの実行と、出力の確認
std::vector<Tensor> outputs;
// 2番目の引数に、出力の欲しいノードの名前を渡す
session->Run({},{"my_ops_out"},{}, &outputs);
// outputsから1番目のTensorを取り出す
auto b = outputs[0].flat<int32>();
for (int i = 0; i < b.size(); ++i)
std::cout << b(i) << ",";
std::cout << std::endl;
実行
$ bazel build tensorflow/my_work/...
$ bazel-bin/tensorflow/my_work/my_work
my op!
6,4,
入力を2倍しただけですが、うまくいくとうれしいですね
その他
TensorFlowのRoadmapをふと眺めていたら、
Improve support for C++ only users
Graph construction
Gradients
Shape Inference
という項目が…
というか先ほどのtensor_c_apiとかのことだろうか？