Class QueueInput

tensorpack/input_source/input_source.py:189–276

Enqueue datapoints from a DataFlow to a TF queue. And the model receives dequeued tensors.

class QueueInput(FeedfreeInput):
    """ Enqueue datapoints from a DataFlow to a TF queue.
        And the model receives dequeued tensors.
    """

    def __init__(self, ds, queue=None):
        """
        Args:
            ds(DataFlow): the input DataFlow.
            queue (tf.QueueBase): A :class:`tf.QueueBase` whose type
                should match the corresponding input signature of the model.
                Defaults to a FIFO queue of size 50.
        """
        if not isinstance(ds, DataFlow):
            raise ValueError("QueueInput takes a DataFlow! Got {}".format(ds))
        self.queue = queue
        self.ds = ds
        self._inf_ds = RepeatedData(ds, -1)
        self._started = False

    def _size(self):
        return len(self.ds)

    def _setup(self, inputs):
        self._input_placehdrs = [build_or_reuse_placeholder(v) for v in inputs]
        assert len(self._input_placehdrs) > 0, \
            "QueueInput has to be used with some inputs!"
        with self.cached_name_scope():
            if self.queue is None:
                self.queue = tfv1.FIFOQueue(
                    50, [x.dtype for x in self._input_placehdrs],
                    name='input_queue')
            logger.info("Setting up the queue '{}' for CPU prefetching ...".format(self.queue.name))
            self.thread = EnqueueThread(self.queue, self._inf_ds, self._input_placehdrs)

            self._dequeue_op = self.queue.dequeue(name='dequeue_for_reset')

    def refill_queue(self):
        """
        Clear the queue, then call dataflow.__iter__() again and fill into the queue.
        """
        self.thread.pause()     # pause enqueue

        opt = tfv1.RunOptions()
        opt.timeout_in_ms = 2000   # 2s
        sess = tfv1.get_default_session()
        # dequeue until empty
        try:
            while True:
                sess.run(self._dequeue_op, options=opt)
        except tf.errors.DeadlineExceededError:
            pass

        # reset dataflow, start thread
        self.thread.reinitialize_dataflow()
        self.thread.resume()

    def _create_ema_callback(self):