Tensorflow: Viết một Op trong Python

Question

Tôi muốn viết một Op bằng Python. Hướng dẫn này chỉ giải thích cách thực hiện nó trong C++ với một trình bao bọc Python. https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html#adding-a-new-op

Tôi có thể viết hoàn toàn bằng Python bằng cách nào?

Answer 1

Bạn có thể sử dụng tf.py_func(func, inp, Tout).

Kết thúc hàm python và sử dụng chức năng này như là một chuỗi lưu lượng.

Cho một hàm python func, có các mảng cố định làm đầu vào của nó và trả về các mảng có nhiều mảng như đầu ra của nó.

---

chức năng python của bạn cần có:

• mảng NumPy như đầu vào, cho ăn từ đồ thị với các đối số inp
• mảng NumPy như kết quả đầu ra, bạn cần phải xác định loại của họ để TensorFlow trong đối số Tout

Bên trong hàm, bạn có thể làm bất cứ điều gì y ou như, nếu điều kiện của vòng lặp, bất cứ điều gì mà không thể có trong TensorFlow.

---

Tuy nhiên, thao tác sẽ được thực thi trên CPU để nó có thể chậm hơn so với TensorFlow op tương đương trong GPU.

Answer 2

Bạn có thể sử dụng tf.py_func để gọi hàm python. Các hoạt động bên trong chức năng cũng có thể trên GPU. Ví dụ, chúng ta có thể thêm Op và gradient của nó hoàn toàn bằng python gọi Caffè trên 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