Class Plan

caffe2/python/core.py:2882–2956

class Plan:

    def __init__(self, name_or_step):
        self._plan = caffe2_pb2.PlanDef()
        self._net_dict = OrderedDict()
        self._steps = []    # A list of ExecutionStep
        if isinstance(name_or_step, ExecutionStep):
            self._plan.name = name_or_step.Name()
            self.AddStep(name_or_step)
        elif isinstance(name_or_step, basestring):
            self._plan.name = name_or_step
        else:
            raise ValueError('name_or_step must be a string or ExecutionStep')

    def __str__(self):
        return self._plan.name

    def Proto(self):
        return self._plan

    def AddNets(self, nets):
        for net in nets:
            if _add_net_to_dict(self._net_dict, net):
                assert isinstance(net, Net)
                self._plan.network.add().CopyFrom(net.Proto())

    def Nets(self):
        return list(self._net_dict.values())

    def AddStep(self, step):
        assert isinstance(step, ExecutionStep)
        step._notify_is_used()
        if not step.HasNets() and not step.HasSubsteps():
            return
        self._plan.execution_step.add().CopyFrom(step.Proto())
        self._steps.append(step)
        # nets need to be added to the plan in order of usage
        net_list = []
        add_nets_in_order(step, net_list)
        self.AddNets([step.get_net(n) for n in net_list])

    def Steps(self):
        return self._steps

    def get_all_attributes(self, name):
        """
        Return the list of all attributes under the given `name`, present in
        all of the nets used in this plan.
        """
        return [
            attr
            for net in self._net_dict.values()
            for attr in net.get_attributes(name)
        ]

    @classmethod
    def create_from_proto(cls, plan_proto):
        assert isinstance(plan_proto, caffe2_pb2.PlanDef)