Class ProgramConfig

test/ir/inference/program_config.py:244–342

A config builder for generating a Program. input_type : (np.dtype, default=None), the inputs will be casted to input_type before fed into TRT engine. If set to None, no casting will be performed. no_cast_list : (list[str], default=None), specify the tensors that will skip the

class ProgramConfig:
    '''A config builder for generating a Program.
    input_type : (np.dtype, default=None), the inputs will be casted to input_type before
                fed into TRT engine. If set to None, no casting will be performed.
    no_cast_list : (list[str], default=None), specify the tensors that will skip the casting
    '''

    def __init__(
        self,
        ops: list[OpConfig],
        weights: dict[str, TensorConfig],
        inputs: dict[str, TensorConfig],
        outputs: list[str],
        input_type: np.dtype | None = None,
        no_cast_list: list[str] | None = None,
    ):
        self.ops = ops
        # if no weight need to save, we create a place_holder to help serialize params.
        if not weights:

            def generate_weight():
                return np.array([1]).astype(np.float32)

            self.weights = {
                "place_holder_weight": TensorConfig(data_gen=generate_weight)
            }
        else:
            self.weights = weights
        self.inputs = inputs
        self.outputs = outputs
        self.input_type = input_type
        self.no_cast_list = [] if no_cast_list is None else no_cast_list
        self.supported_cast_type = [np.float32, np.float16]

    def __repr__(self):
        log_str = ''
        for i in range(len(self.ops)):
            if i != len(self.ops) - 1:
                log_str += repr(self.ops[i]) + ' + '
            else:
                log_str += repr(self.ops[i])
        log_str += ' -- '
        for t, v in self.inputs.items():
            log_str += '[' + t + ': ' + str(v) + ']'
        for t, v in self.weights.items():
            log_str += '[' + t + ': ' + str(v) + ']'
        log_str += f"['input_type': {self.input_type}]"
        return log_str

    def set_input_type(self, _type: np.dtype) -> None:
        assert _type in self.supported_cast_type or _type is None, (
            "PaddleTRT only supports FP32 / FP16 IO"
        )

        ver = paddle.inference.get_trt_compile_version()
        trt_version = ver[0] * 1000 + ver[1] * 100 + ver[2] * 10
        if trt_version < 8600:
            logging.info("set_input_type is ignored for TRT version < 8600")