Method infer_shapes

tensorrt_llm/runtime/session.py:204–237

@brief: Set input shapes to given context, and infer the output shapes from the given input shapes. This function should be called every time when the input shapes are changed before calling run(). Or call the context.set_input_shape on all dynamic shaped input

(
            self,
            inputs: List[TensorInfo],
            context: Optional[trt.IExecutionContext] = None
    )

                    )

    def infer_shapes(
            self,
            inputs: List[TensorInfo],
            context: Optional[trt.IExecutionContext] = None
    ) -> List[TensorInfo]:
        '''
        @brief: Set input shapes to given context, and infer the output shapes from the given input shapes.
               This function should be called every time when the input shapes are changed before calling run().
               Or call the context.set_input_shape on all dynamic shaped input tensors manually.
        @param inputs: list of TensorInfo object, each item represents an input tensor
        @param context: TensorRT execution context, if None, use the default context
        @return: list of TensorInfo object, each item represents an output tensor, returns None if failed
        '''
        # set shape to the default context if context is not specified
        if context is None:
            context = self.context
        for i in inputs:
            if self.engine.get_tensor_mode(i.name) != trt.TensorIOMode.INPUT:
                raise ValueError(f"Tensor:{i.name} is not an input tensor")
            if self.engine.get_tensor_dtype(i.name) != i.dtype:
                raise ValueError(f"Tensor:{i.name} has wrong dtype")
            if not context.set_input_shape(i.name, i.shape):
                raise RuntimeError(
                    f"Could not set shape {i.shape} for tensor {i.name}. Please check the profile range for which your model was build."
                )

        outputs = []
        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
                shape = context.get_tensor_shape(name)
                dtype = self.engine.get_tensor_dtype(name)
                outputs.append(TensorInfo(name, dtype, shape))
        return outputs

    def _set_weight_streaming(self, gpu_weights_percent):
        if not self.engine.streamable_weights_size: