Awesome
TensorFlow ZMQ Op
A convenient way to receive data from other processes. This small library can:
- Send a list of numpy arrays from python; serialization is written in C++ for efficiency.
- One copy in merging all the buffers; One copy in pybind11 overhead (TODO); One copy in ZMQ send.
- Receive a list of tensors from tensorflow;
- One copy in ZMQ recv; One copy to split the buffer into tensors.
- The op is stateful and safe to be evaluated multiple times in one
sess.run
call.
- Serialization is in a custom protocol for efficiency;
Why:
Sometimes for complicated large-scale tasks you would really want data processing to be separate from TensorFlow. However in TensorFlow there is no good way to receive data from other processes.
Build:
Require gcc>=5.3, tensorflow>=1.4, zeromq>=4.
Require the zmq.hpp
header from cppzmq at
your compiler's include path, or under the src
directory.
Add /path/to/git/clone/zmq_ops
to PYTHONPATH
to be able to import it.
Or use pip install .
to install it.
Ops will be compiled the first time it gets imported. Note that it usually requires recompilation after a TensorFlow reinstallation.
Use:
See benchmark.py
for usage.
On my machine this script can achieve about 1.3GB/s throughput. Equivalent to about 2.3k float32 (or 9.2k uint8) imagenet images per second.