Tensorflow: Ein Op in Python schreiben

Question

Ich würde gerne ein Op in Python schreiben. Dieses Tutorial erklärt nur, wie man es in C++ mit einem Python-Wrapper macht. https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html#adding-a-new-op

Wie kann ich es vollständig in Python schreiben?

Answer 1

Sie können tf.py_func(func, inp, Tout) verwenden.

Wraps eine Python-Funktion und verwendet sie als Tensorflow-Op.

Gegeben eine Python-Funktion func, die numpy-Arrays als Eingaben akzeptiert und numpy-Arrays als Ausgaben zurückgibt.

---

Ihre Python-Funktion muss haben:

• numpy Arrays als Eingaben aus dem Diagramm mit dem Argument zugeführt inp
• numpy Arrays als Ausgänge, Sie brauchen, um ihre Typen zu spezifizieren TensorFlow im Argument Tout

Innerhalb der Funktion können Sie tun, was y Ou, wenn Bedingungen für for-Schleifen, alles, was in TensorFlow nicht möglich ist.

---

Jedoch wird der Betrieb auf der CPU ausgeführt werden, so dass er sich langsamer als die äquivalente TensorFlow op in GPU sein kann.

Answer 2

Sie können tf.py_func verwenden, um Python-Funktionen aufzurufen. Die Operationen innerhalb der Funktion können auch auf der GPU ausgeführt werden. Zum Beispiel können wir einen Op und seinen Gradienten rein in Python hinzufügen, der Caffe auf GPU aufruft:

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