組み込みTensorFlowカーネルを変更する最も良い方法

Question

TensorFlow組み込み演算子カーネルを変更する最も良い方法を学びたいと思います。 たとえば、static const double Aの値をtensorflow/core/kernels/resize_bicubic_op.ccに変更したいとします。

1. TensorFlowライブラリ全体を直接変更して再コンパイルします。この解法の問題点は次のとおりです。A.双三次補間を使用するすべての関数に影響します。 B.これはライブラリ全体を再コンパイルする必要があり、バイナリからインストールするときには動作しません。

2. カスタムオペレーションとして定義します。問題は、ソースコード内にREGISTER_OP()がないことです。私はこの双三次関数のためにREGISTER_OP()を書く方法と、他の修正が必要かどうか分からない。

他にも良い方法がありますか？

ありがとうございました。

Answer 1

この問題にアプローチする最善の方法は、カスタム操作を作成することです。カスタム操作を追加する方法の詳細については、this tutorialを参照してください。 REGISTER_OPは、tf.image.resize_bicubic() opの呼び出しはtensorflow/core/ops/image_ops.ccです。

別の方法としては、再使用するのと同じOP登録することであり、代替実装で新しいカーネルを登録します。これにより、（実験的）Graph.kernel_label_map() APIを使用して"ResizeBicubic" opの代替実装を選択することができます。

_ = tf.load_op_library(...) # Load the .so containing your implementation. 

with tf.get_default_graph().kernel_label_map({"ResizeBicubic": "my_impl"}): 
    images = tf.image.resize_bicubic(...) # Will use your implementation. 

を...そして、あなたのC++コードでラベル"my_impl"を指定するカーネルの登録を追加します：たとえば、あなたはPythonプログラムで次の操作を行うことができ

template <typename Device, typename T> 
class MyResizeBicubicOp<Device, T> : public OpKernel { 
    // Custom implementation goes here... 
} 

#define REGISTER_KERNEL(T)       \ 
    REGISTER_KERNEL_BUILDER(Name("ResizeBicubic")  \ 
           .Device(DEVICE_CPU)  \ 
           .Label("my_impl")  \ 
           .TypeConstraint<T>("T") \ 
           .HostMemory("size"), \ 
          MyResizeBicubicOp<CPUDevice, T>); 

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);