Class ProcessingReader

caffe2/python/pipeline.py:353–394

Reader that reads from an upstream reader, calls the processor, and returns the processed record.

class ProcessingReader(Reader):
    """
    Reader that reads from an upstream reader, calls the processor, and returns
    the processed record.
    """
    def __init__(self, reader, processor):
        Reader.__init__(self)
        self.reader = reader
        self.processor = make_processor(processor, reader)

    def schema(self):
        return self.processor.schema()

    def setup_ex(self, init_net, finish_net):
        self.reader.setup_ex(init_net, finish_net)

    def read_ex(self, init_net, exit_net):
        read_nets, status, rec = self.reader.read_record_ex(init_net, exit_net)
        # We don't use status as stop_blob of NetBuilder it's not guarantee that
        # it would end up being the true stob_blob. For example,
        # ReaderWithLimitBase doesn't pass the status through but rather copy
        # from it.
        with NetBuilder() as nb:
            # Current NetBuilder is optionally used inside the processor,
            # then its children are retrieved inside of
            # normalize_processor_output.
            # Once readers and writers also use NetBuilder,
            # this logic will be more natural.
            result = normalize_processor_output(self.processor(rec))
        read_nets += result.nets
        if result.should_stop or nb._stop_blob:
            stop_net = core.Net('stop_net')
            if result.should_stop:
                stop_net.Or([status, result.should_stop], [status])
            if nb._stop_blob:
                stop_net.Or([status, nb._stop_blob], [status])
            read_nets.append(stop_net)
        if hasattr(self.processor, 'setup'):
            init_net.add_attribute(TaskGroup.LOCAL_SETUP, self.processor)
        self._set_schema(result.record)
        fields = result.record.field_blobs() if result.record else None
        return read_nets, status, fields


class NetProcessor: