Tensorflow：PythonでOpを書く

Question

OpをPythonで書いてみたいと思います。このチュートリアルでは、Pythonラッパーを使用してC++で行う方法についてのみ説明します。 https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html#adding-a-new-op

どうすれば完全にPythonで記述できますか？

Answer 1

tf.py_func(func, inp, Tout)を使用できます。

python関数をラップし、テンソルフロー演算子として使用します。

numpy配列を入力として受け取り、numpy配列をその出力として返す、python関数funcを指定します。

---

あなたのPythonの機能は持っている必要があります：

• numpyの配列を引数に出力としてinp
• numpyの配列をグラフから供給された入力は、あなたが自分のタイプを指定する必要があるとして、引数の中でTensorFlowを使う

関数の中では、yもしforループの条件なら、TensorFlowでは不可能なことがあります。

---

ただし、操作はCPU上で実行されるため、GPUの同等のTensorFlow操作よりも遅くなる可能性があります。

Answer 2

tf.py_funcを使用すると、python関数を呼び出すことができます。関数内の操作は、GPU上で行うこともできます。

def custom_loss_impl(x): 
    caffe.set_mode_gpu() 
    caffe.set_device(0) 
    ... 
    return np.float32(loss) 

def custom_loss(x): 
    tf.RegisterGradient("custom_loss_grad")(custom_loss_grad) 
    g=tf.get_default_graph() 
    with g.gradient_override_map({"PyFunc":"custom_loss_grad"}): 
     return tf.py_func(custom_loss_impl,[x],[tf.float32])[0] 

def custom_loss_grad_impl(x): 
    caffe.set_mode_gpu() 
    caffe.set_device(0) 
    custom_loss_impl(x) 
    ... 
    return np.float32(gradient) 

def custom_loss_grad(op,grad): 
    x=op.inputs[0] 
    return tf.py_func(custom_loss_grad_impl,[x],[tf.float32])#assume grad=1