Long Inference Time on First Run After Changing Input Shape in Dynamic Shape TensorRT Engine

Description

I have modified the inference part of the code from the Stable Diffusion demo and generated a TensorRT engine with dynamic shapes. However, every time I change the input shape, the inference time is very long, which makes it difficult to perform inference with user-defined image sizes. Even small changes in input shape lead to significant variations in inference time.

I am running inference on the Unet+Controlnet part of the Stable Diffusion demo, where both the input and output have dynamic shapes. The code I modified is based on the following implementation, and I am not using CUDA Graph during inference.

The range of shapes I am working with is quite large, from a minimum of [1, 3, 1, 1] to a maximum of [8, 3, 1280, 1280]. However, when I adjust the input shape from [2, 3, 1024, 1024] to [2, 3, 1024, 960], the inference time for the first run is very long after the shape changed, and only on the second run does the inference time meet expectations. So that the feature for custom image sizes cannot be implement.

Could you help me understand why this behavior occurs and how to optimize the inference time for changing input and output shapes?

Code

    def infer_with_dynamic_shape(self, feed_dict, shape_dict, stream, use_cuda_graph=False):

		for binding in range(self.engine.num_io_tensors):
            tensor_name = self.engine.get_tensor_name(binding)
            if self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
                input_tensor = feed_dict[tensor_name]
                if not isinstance(input_tensor, torch.Tensor) or self.tensors[tensor_name].shape == input_tensor.shape and self.tensors[tensor_name].device == input_tensor.device:
                    self.tensors[tensor_name].copy_(input_tensor)
                else:
                    device = self.tensors[tensor_name].device
                    dtype = self.tensors[tensor_name].dtype
                    self.tensors[tensor_name] = torch.empty(shape_dict[tensor_name], dtype=dtype).to(device=device)
                    self.tensors[tensor_name].copy_(input_tensor)
                    self.context.set_input_shape(tensor_name, shape_dict[tensor_name])
            else:
                if shape_dict[tensor_name] and shape_dict[tensor_name] != self.tensors[tensor_name].shape:
                    device = self.tensors[tensor_name].device
                    dtype = self.tensors[tensor_name].dtype
                    pre_shape = self.tensors[tensor_name].shape
                    self.tensors[tensor_name] = torch.empty(shape_dict[tensor_name], dtype=dtype).to(device=device)

        for name, tensor in self.tensors.items():
            self.context.set_tensor_address(name, tensor.data_ptr())

        if use_cuda_graph:
            if self.cuda_graph_instance is not None:
                CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream))
                CUASSERT(cudart.cudaStreamSynchronize(stream))
            else:
                # do inference before CUDA graph capture
                noerror = self.context.execute_async_v3(stream)
                if not noerror:
                    raise ValueError(f"ERROR: inference failed.")
                # capture cuda graph
                CUASSERT(
                    cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal))
                self.context.execute_async_v3(stream)
                self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream))
                self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
        else:
            noerror = self.context.execute_async_v3(stream)
            if not noerror:
                cuda_error = cudart.cudaGetLastError()
                raise ValueError(f"ERROR: inference failed. ")

Environment

TensorRT Version: 10.7.0

NVIDIA GPU: Nvidia L20

NVIDIA Driver Version: 535.216.01

CUDA Version: 12.2

CUDNN Version: