Class CompositeReader

caffe2/python/dataio.py:519–573

Base class for a reader that wrap multiple readers, e.g., reading from multiple sources simultaneously.

class CompositeReader(Reader):
    """
    Base class for a reader that wrap multiple readers, e.g., reading from
    multiple sources simultaneously.
    """
    def __init__(self, names, readers):
        """
        Args:
            names: list[str] names of readers; used as schema keys
            readers: list[Reader] Reader instances, must have schema
        """
        assert len(names) == len(readers)
        super().__init__(schema=Struct(*[
            (name, reader.schema()) for name, reader in zip(names, readers)
        ]))
        self._names = names
        self._readers = readers

    def setup_ex(self, init_net, finish_net):
        for reader in self._readers:
            reader.setup_ex(init_net, finish_net)

    def read_ex(self, local_init_net, local_finish_net):
        """
        Stops when one of the reader finished
        """
        # First, instantiate all the reader nets
        fields = []
        stop_blobs = []
        all_sub_read_nets = []
        for name, reader in zip(self._names, self._readers):
            sub_read_nets, should_stop, record = reader.read_record_ex(
                local_init_net, local_finish_net)
            stop_blobs.append(should_stop)
            all_sub_read_nets.append(sub_read_nets)
            fields.extend(record.field_blobs())

        read_nets = []
        # Use the stop blob of the last reader as stop blob of composite reader.
        local_should_stop = stop_blobs[-1]
        for name, sub_read_nets, stop_blob in zip(self._names, all_sub_read_nets, stop_blobs):
            read_nets.extend(sub_read_nets)
            if stop_blob == local_should_stop:
                # Skip adding stop net because Or([A, A], A) doesn't pass operator
                # schema check
                continue
            stop_net = core.Net("{}_stop".format(name))
            stop_net.Or([local_should_stop, stop_blob], local_should_stop)
            read_nets.append(stop_net)

        return read_nets, local_should_stop, fields

    def reset(self, net):
        for reader in self._readers:
            reader.reset(net)


class CompositeReaderBuilder(ReaderBuilder):