MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / _add_optimization_profile

Method _add_optimization_profile

tensorrt_llm/builder.py:238–274  ·  view source on GitHub ↗
(self, network: Network,
                                  builder_config: BuilderConfig)

Source from the content-addressed store, hash-verified

236 **kwargs)
237
238 def _add_optimization_profile(self, network: Network,
239 builder_config: BuilderConfig):
240 assert isinstance(builder_config, BuilderConfig)
241 assert isinstance(network, Network)
242 input_tensors = network._inputs
243 if len(input_tensors) == 0:
244 logger.warning("There are no inputs in the network!")
245 return
246 num_profiles = len(list(input_tensors.values())[0].profiles)
247 force_num_profiles = getattr(builder_config, "force_num_profiles", None)
248 for i in range(num_profiles):
249 logger.debug(f'Adding optimization profile {i+1}/{num_profiles}')
250 profile = self.trt_builder.create_optimization_profile()
251 for input_name in input_tensors.keys():
252 if len(input_tensors[input_name].profiles) == 0:
253 continue
254 shape_profile = input_tensors[input_name].profiles[i]
255 min_shape = [*shape_profile.min]
256 opt_shape = [*shape_profile.opt]
257 max_shape = [*shape_profile.max]
258 profile.set_shape(input_name, min_shape, opt_shape, max_shape)
259 logger.debug(
260 f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}, dimension names: {shape_profile.dimension_names}'
261 )
262 ret = builder_config.trt_builder_config.add_optimization_profile(
263 profile)
264 logger.debug(f"Added optimization profile: #{ret}")
265 if force_num_profiles is not None and (
266 i + 1
267 ) == force_num_profiles and force_num_profiles < num_profiles:
268 logger.warning(
269 f"Only adding {force_num_profiles} profiles instead of {num_profiles}."
270 )
271 break
272 assert self._validate_named_dimensions(
273 network, builder_config
274 ), "Validation of the tensor dimension ranges failed, please check the dimension ranges, find the offensive tensor and dimension name in above the error log"
275
276 def _validate_named_dimensions(self, network: Network,
277 builder_config) -> bool:

Callers 1

build_engineMethod · 0.95

Calls 5

warningMethod · 0.45
valuesMethod · 0.45
debugMethod · 0.45
keysMethod · 0.45

Tested by

no test coverage detected