diff --git a/tools/Polygraphy/CHANGELOG.md b/tools/Polygraphy/CHANGELOG.md index d6da60ea3..957750031 100644 --- a/tools/Polygraphy/CHANGELOG.md +++ b/tools/Polygraphy/CHANGELOG.md @@ -13,6 +13,8 @@ Dates are in YYYY-MM-DD format. ### Fixed - Fixed issue when `polygraphy multi-device shard` would exceed python recursive depth limit on large models. +- Fixed a bug where `polygraphy convert` would fail when the input model was already a + TensorRT engine, e.g. `polygraphy convert model.engine --convert-to trt -o out.plan`. - Fixed a bug where setting `POLYGRAPHY_ASK_BEFORE_INSTALL=0` would enable asking before installing dependencies instead of disabling it. diff --git a/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py b/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py index e78b8fa38..87cf7db9a 100644 --- a/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py +++ b/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py @@ -712,14 +712,14 @@ def add_to_script_impl(self, script, network_name=None): def load_engine_bytes(self, network=None): """ - Loads a TensorRT engine according to arguments provided on the command-line. + Loads serialized TensorRT engine bytes according to arguments provided on the command-line. Args: network (Tuple[trt.Builder, trt.INetworkDefinition, Optional[parser]]): A tuple containing a TensorRT builder, network and optionally parser. Returns: - tensorrt.ICudaEngine: The engine. + bytes: The serialized engine. """ loader = args_util.run_script(self.add_to_script, network) return loader() diff --git a/tools/Polygraphy/polygraphy/tools/convert/convert.py b/tools/Polygraphy/polygraphy/tools/convert/convert.py index 50ba59bb5..6fb7e3bbf 100644 --- a/tools/Polygraphy/polygraphy/tools/convert/convert.py +++ b/tools/Polygraphy/polygraphy/tools/convert/convert.py @@ -101,11 +101,11 @@ def run_impl(self, args): model = self.arg_groups[OnnxLoadArgs].load_onnx() self.arg_groups[OnnxSaveArgs].save_onnx(model, args.output) elif convert_type.is_trt(): - with self.arg_groups[ + serialized_engine = self.arg_groups[ TrtLoadEngineBytesArgs - ].load_engine_bytes() as serialized_engine: - self.arg_groups[TrtSaveEngineBytesArgs].save_engine_bytes( - serialized_engine, args.output - ) + ].load_engine_bytes() + self.arg_groups[TrtSaveEngineBytesArgs].save_engine_bytes( + serialized_engine, args.output + ) else: G_LOGGER.critical(f"Cannot convert to model type: {convert_type}") diff --git a/tools/Polygraphy/tests/tools/test_convert.py b/tools/Polygraphy/tests/tools/test_convert.py index 19f313deb..5739bf8a8 100644 --- a/tools/Polygraphy/tests/tools/test_convert.py +++ b/tools/Polygraphy/tests/tools/test_convert.py @@ -145,6 +145,19 @@ def test_modify_onnx_outputs(self, poly_convert): model = onnx.load(outmodel.name) assert len(model.graph.output) == 2 + def test_engine_to_engine(self, poly_convert): + # Engine inputs yield plain bytes, not an object that supports the + # context manager protocol. + engine_bytes = b"serialized-engine" + with util.NamedTemporaryFile( + "w+b", suffix=".engine" + ) as inmodel, util.NamedTemporaryFile(suffix=".plan") as outmodel: + inmodel.write(engine_bytes) + inmodel.flush() + + poly_convert([inmodel.name, "--model-type=engine", "-o", outmodel.name]) + assert BytesFromPath(outmodel.name)() == engine_bytes + class TestConvertToOnnxLikeTrt: @pytest.mark.parametrize(