def __call__(self): """ Builds a TensorRT engine. Returns: trt.ICudaEngine: The engine that was created. """ # If network is a callable, then we own its return value ret, owning = misc.try_call(self._network) builder, network, parser = misc.unpack_args(ret, num=3) with contextlib.ExitStack() as stack: provided = "Builder and Network" if parser is None else "Builder, Network, and Parser" if owning: stack.enter_context(builder) stack.enter_context(network) if parser is not None: stack.enter_context(parser) else: G_LOGGER.verbose("{:} were provided directly instead of via a Callable. This loader will not assume ownership. " "Please ensure that they are freed.".format(provided)) network_log_mode = "full" if G_LOGGER.severity <= G_LOGGER.ULTRA_VERBOSE else "attrs" G_LOGGER.super_verbose(lambda: ("Displaying TensorRT Network:\n" + trt_util.str_from_network(network, mode=network_log_mode))) config, _ = misc.try_call(self._config, builder, network) G_LOGGER.info("Building engine with configuration: {:}".format(trt_util.str_from_config(config))) engine = builder.build_engine(network, config) if not engine: G_LOGGER.critical("Invalid Engine. Please ensure the engine was built correctly") if hasattr(config.int8_calibrator, "free"): config.int8_calibrator.free() return engine
def call_impl(self): """ Returns: bytes: The serialized engine that was created. """ # If network is a callable, then we own its return value ret, owns_network = util.invoke_if_callable(self._network) builder, network, parser = util.unpack_args(ret, num=3) if builder is None or network is None: G_LOGGER.critical("Expected to recevie a (builder, network) tuple for the `network` parameter, " "but received: ({:}, {:})".format(builder, network)) with contextlib.ExitStack() as stack: if owns_network: stack.enter_context(builder) stack.enter_context(network) if parser is not None: stack.enter_context(parser) else: provided = "Builder and Network" if parser is None else "Builder, Network, and Parser" G_LOGGER.verbose("{:} were provided directly instead of via a Callable. This loader will not assume ownership. " "Please ensure that they are freed.".format(provided)) config, owns_config = util.invoke_if_callable(self._config, builder, network) if owns_config: stack.enter_context(config) else: G_LOGGER.verbose("Builder configuration was provided directly instead of via a Callable. This loader will not assume " "ownership. Please ensure it is freed.") try: config.int8_calibrator.__enter__ # Polygraphy calibrator frees device buffers on exit. except AttributeError: pass else: stack.enter_context(config.int8_calibrator) network_log_mode = "full" if G_LOGGER.severity <= G_LOGGER.ULTRA_VERBOSE else "attrs" G_LOGGER.super_verbose(lambda: ("Displaying TensorRT Network:\n" + trt_util.str_from_network(network, mode=network_log_mode))) G_LOGGER.start("Building engine with configuration:\n{:}".format(trt_util.str_from_config(config))) try: engine_bytes = builder.build_serialized_network(network, config) except AttributeError: engine = builder.build_engine(network, config) if not engine: G_LOGGER.critical("Invalid Engine. Please ensure the engine was built correctly") stack.enter_context(engine) engine_bytes = engine.serialize() if not engine_bytes: G_LOGGER.critical("Invalid Engine. Please ensure the engine_bytes was built correctly") try: timing_cache = config.get_timing_cache() except AttributeError: if self.timing_cache_path: trt_util.fail_unavailable("save_timing_cache in EngineBytesFromNetwork") else: if timing_cache and self.timing_cache_path: with timing_cache.serialize() as buffer: util.save_file(buffer, self.timing_cache_path, description="tactic timing cache") return engine_bytes