def _setup(self, inputs_desc):
assert len(inputs_desc) > 0, \
"ZMQInput has to be used with InputDesc!"
self._desc = inputs_desc
import zmq_ops
self._zmq_pull_socket = zmq_ops.ZMQPullSocket(
self._end_point, [x.type for x in inputs_desc], self._hwm)
def _setup(self, input_signature):
assert len(input_signature) > 0, \
"ZMQInput has to be used with input signature!"
import zmq_ops
self._zmq_pull_socket = zmq_ops.ZMQPullSocket(
self._end_point, [x.dtype for x in input_signature],
hwm=self._hwm,
bind=self._bind)
def to_dataset(self, input_signature):
"""
Convert to a TF dataset.
Args:
input_signature (list[InputSpec]):
Returns:
tf.data.Dataset
"""
import zmq_ops
zmq_pull_socket = zmq_ops.ZMQPullSocket(
self._end_point, [x.dtype for x in input_signature],
hwm=self._hwm, bind=self._bind)
def mapper(_):
inputs = list(zmq_pull_socket.pull())
for v, sig in zip(inputs, input_signature):
v.set_shape(sig.shape)
return inputs
# Is there a better way to construct from stateful tensor?
dataset = tf.data.Dataset.from_tensors([1]) # just a placeholder
return dataset.map(mapper)
