Class RunnableGraph

ppq/IR/deploy.py:13–136

class RunnableGraph(GraphCommandProcessor):
    def __init__(self, graph: BaseGraph, device: str = None):
        """RunnableGraph deals with values related with graph executing.

        Literally it helps you move values of your graph towards device and vice versa.
            And give an executable order of all operations in your graph which actual executor will follow.
        Args:
            graph (BaseGraph): BaseGraph instance.
            device (str, optional): This attribute is only used by with RunnableGraph(graph, device) syntactic.
            next_command_processor (Callable, optional): next processor in processing chain.
        """
        super().__init__(graph_or_processor=graph)
        self._device = device  # only used in "with RunnableGraph(graph, device):"

    def process(self, command: GraphCommand) -> Any:

        if command.command_type == GraphCommandType.DEPLOY_TO_CPU:
            return self.deploy('cpu')

        elif command.command_type == GraphCommandType.DEPLOY_TO_CUDA:
            if isinstance(command, GraphDeployCommand):
                device = command._device
                return self.deploy(device)
            else:
                return self.deploy('cuda')

        elif command.command_type == GraphCommandType.DEPLOY_TO_NUMPY:
            return self.retrieve()

    def __enter__(self):
        self.deploy(self._device)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.retrieve()

    def _acceptable_command_types(self) -> List[GraphCommandType]:
        return [
            GraphCommandType.DEPLOY_TO_CPU,
            GraphCommandType.DEPLOY_TO_CUDA,
            GraphCommandType.DEPLOY_TO_NUMPY
        ]

    def retrieve(self):

        for _, operator in self._graph.operations.items():

            assert isinstance(operator, Operation), \
                f'Failed to retrieve graph to numpy, incorrect operator {operator} found.'

            # in onnx format, some constant values are warpped with operation's attributes['value']
            # To move those constant value from numpy to device,
            # we have to move all the attributes['value'] of operation to device(if there is any).
            if operator.type == 'Constant':
                operator.attributes['value'] = \
                    convert_any_to_numpy(operator.attributes['value'])

        for _, variable in self._graph.variables.items():
            assert isinstance(variable, Variable), \