Class ReaderWithLimit

caffe2/python/dataio.py:410–445

Reader that stops after `num_iter` batches. If `num_iter` <= 0 or is None, reverts to an unconstrained reader that exports a boolean blob indicating that the reader has exhausted the data steam.

class ReaderWithLimit(ReaderWithLimitBase):
    """Reader that stops after `num_iter` batches.

    If `num_iter` <= 0 or is None, reverts to an unconstrained reader that
    exports a boolean blob indicating that the reader has exhausted
    the data steam.
    """
    def __init__(self, reader, num_iter=1):
        """Class initializer.

        Args:
            reader: The underlying reader object doing the actual read.
            num_iter: Number of batches to read. If `None`,
                the class reverts to a normal reader except that it also
                produces a data_finished blob as a side effect to indicate
                whether the input stream is exhausted.
        """
        super().__init__(reader)
        self.counter = None
        self.num_iter = num_iter
        if self.num_iter is not None:
            self.counter = self.net.AddExternalInput(
                self.net.NextName('counter'))

    def setup_limiter(self, global_init_net, global_finish_net):
        if self.counter:
            global_init_net.CreateCounter(
                [], [self.counter], init_count=int(self.num_iter))

    def check_limiter_condition(self, stop_condition_net):
        if self.counter:
            return stop_condition_net.CountDown([self.counter], 1)
        else:
            return stop_condition_net.ConstantFill(
                [], 1,
                shape=[], value=False, dtype=core.DataType.BOOL)


def CountUntil(num_iter):