Class QuantableGraph

ppq/IR/quantize.py:259–324

class QuantableGraph(GraphCommandProcessor):
    def process(self, command: GraphCommand) -> Any:
        if command.command_type == GraphCommandType.QUANTIZE_OPERATION:
            assert isinstance(command, QuantizeOperationCommand)
            return self.quantize_operation(
                command.op_name, command.target_platform, command.config)

    def _acceptable_command_types(self) -> List[GraphCommandType]:
        return [
            GraphCommandType.QUANTIZE_OPERATION,
        ]

    def quantize_operation(
        self,
        operation_name: str,
        target_platform: TargetPlatform,
        quantization_config: OperationQuantizationConfig
    ) -> QuantableOperation:
        if operation_name not in self.graph.operations:
            raise KeyError(f'Operation {operation_name} is not in your graph, Please check your input.')

        operation = self._graph.operations[operation_name]
        quantized_operation = QuantableOperation(
            convert_from=operation,
            quantize_config=quantization_config,
            platform=target_platform,
        )

        # calling other chain responder to replace operation with quantized one.
        if self._next_command_processor is None:
            raise RuntimeError(
                'To replace a operation, your processor chain must have a GraphReplacer Processor.')
        self._next_command_processor(ReplaceOperationCommand(operation_name, quantized_operation))

        # replace all related variable with quantable one.
        for var in quantized_operation.inputs + quantized_operation.outputs:
            if isinstance(var, QuantableVariable): continue
            self._next_command_processor(
                ReplaceVariableCommand(
                    var_name=var.name,
                    replace_to=QuantableVariable(convert_from=var)
                )
            )
        quantized_operation.store_parameter_value()

    def dequantize_operation(
        self,
        operation_name: str
    ) -> Operation:
        if operation_name not in self.graph.operations:
            raise KeyError(f'Operation {operation_name} is not in your graph, Please check your input.')
        operation = self._graph.operations[operation_name]
        if not isinstance(operation, QuantableOperation): return operation
        else: return operation.dequantize()

    def dequantize_graph(self, expire_device: str = 'cpu'):
        """一个方便懒人的函数."""
        for operation in self.graph.operations.values():