py_func just for CPU operations?

Question

According to the docs, I can use tf.py_func to define my own operation. It will receive Numpy arrays and must return Numpy arrays. It's not explicitly written in the doc but I conclude from that that it will be bound to the CPU device?

If I combine this with other ops which could potentially run on GPU, will TF greedily move as much computation to the GPU and automatically transfer the memory between GPU and CPU for my tf.py_func op? (Just like Theano would do.)

Is there anything similar like tf.py_func to define a GPU operation?

Answer 1

py_func is an odd one -- it runs the Python code in the same Python interpreter as the interpreter used to create the op. So if your surrounding ops are GPU there will be GPU<->CPU transfers. Furthermore, there will be a copy to move data between Python address space and TensorFlow address space (ie, memcpy here)

If you are asking how to run Python code on GPU, there's Numba. If you are asking how to have a "function"-like thing in TensorFlow GPU, there's Defun. Generally, TensorFlow is moving in the direction of defining enough numpy features so that you won't need to use numpy, but can implement your functions with native TF primitives instead.

There's another potential way of staying in Python-land while relying on TensorFlow engine. You could make a Python wrapper that implements emulates Python numeric type interface, but delegates actual work to underlying TensorFlow engine. Kind of like how numpy implements that interface and delegates underlying work to BLAS library.

With some tricks you could make this work where the data always stays on GPU. A proof of concept implementation is overviewed here