polygraphy 0.50.3__py3-none-any.whl

polygraphy/tools/debug/subtool/precision.py

@@ -0,0 +1,281 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import math
+
+from polygraphy import mod
+from polygraphy.mod.trt_importer import lazy_import_trt
+from polygraphy.logger import G_LOGGER
+from polygraphy.tools.args import ModelArgs, TrtConfigArgs
+from polygraphy.tools.debug.subtool.base import BaseCheckerSubtool
+
+trt = lazy_import_trt()
+trt_util = mod.lazy_import("polygraphy.backend.trt.util")
+
+
+class BaseMarker:
+    def __init__(self, max_layers, direction, num_layers_to_mark):
+        self.max_layers = max_layers
+        self.direction = direction
+        self.num_layers_to_mark = num_layers_to_mark
+        self.good = (
+            max_layers + 1
+        )  # Pretend marking all the layers gives us good accuracy.
+
+    def select_layers(self):
+        if self.direction == "forward":
+            G_LOGGER.info(
+                f"Selecting first {self.num_layers_to_mark} layer(s) to run in higher precision"
+            )
+            return range(0, self.num_layers_to_mark)
+        else:
+            G_LOGGER.info(
+                f"Selecting last {self.num_layers_to_mark} layer(s) to run in higher precision"
+            )
+            return range(self.max_layers - self.num_layers_to_mark, self.max_layers)
+
+    def success_message(self):
+        which_layers = "first" if self.direction == "forward" else "last"
+        G_LOGGER.finish(
+            f"To achieve acceptable accuracy, try running the {which_layers} {self.good} layer(s) in higher precision"
+        )
+
+
+class BisectMarker(BaseMarker):
+    def __init__(self, max_layers, direction) -> None:
+        super().__init__(max_layers, direction, max_layers)
+        self.bad = 0
+
+    def step(self, success):
+        if success:
+            self.good = self.num_layers_to_mark
+            # On successes, we want num_layers_to_mark to go closer to self.bad
+            round_func = math.floor
+        else:
+            self.bad = self.num_layers_to_mark
+            round_func = math.ceil
+
+        old_num_layers_to_mark = self.num_layers_to_mark
+        self.num_layers_to_mark = int(round_func((self.good + self.bad) / 2.0))
+
+        # Prevent infinite looping:
+        if old_num_layers_to_mark == self.num_layers_to_mark:
+            return True
+
+        # If good and bad are within 1 layer of each other,
+        # then we already have the information we need.
+        if abs(self.good - self.bad) <= 1:
+            if self.good >= self.max_layers:
+                G_LOGGER.error(
+                    "Could not find a configuration that satisfied accuracy requirements."
+                )
+            else:
+                self.success_message()
+            return True
+
+        if self.num_layers_to_mark > self.max_layers:
+            G_LOGGER.error(
+                "Could not find a configuration that satisfied accuracy requirements."
+            )
+            return True
+
+        return False
+
+    def remaining(self):
+        return int(math.ceil(math.log2(self.good - self.bad)))
+
+
+class LinearMarker(BaseMarker):
+    def __init__(self, max_layers, direction) -> None:
+        super().__init__(max_layers, direction, 0)
+
+    def step(self, success):
+        if success:
+            self.good = self.num_layers_to_mark
+        self.num_layers_to_mark += 1
+
+        if success:
+            self.success_message()
+            return True
+
+        if self.num_layers_to_mark > self.max_layers:
+            G_LOGGER.error(
+                "Could not find a configuration that satisfied accuracy requirements."
+            )
+            return True
+
+        return False
+
+    def remaining(self):
+        return self.max_layers - self.num_layers_to_mark
+
+
+class Precision(BaseCheckerSubtool):
+    """
+    [EXPERIMENTAL] Iteratively mark layers to run in a higher precision to find a compromise between performance and quality.
+
+    `debug precision` follows the same general process as other `debug` subtools (refer to the help output
+    of the `debug` tool for more background information and details).
+
+    Each iteration will generate an engine called 'polygraphy_debug.engine' in the current directory.
+    """
+
+    def __init__(self):
+        super().__init__(
+            "precision",
+            precision_constraints_default="obey",
+            allow_no_artifacts_warning=False,
+        )
+
+    def add_parser_args(self, parser):
+        parser.add_argument(
+            "--mode",
+            help="How layers are selected to run in higher precision. "
+            "'bisect' will use binary search, and 'linear' will iteratively mark one extra layer at a time",
+            choices=["bisect", "linear"],
+            default="bisect",
+        )
+        parser.add_argument(
+            "--dir",
+            "--direction",
+            help="Order in which layers are marked to run in higher precision. "
+            "'forward' will start marking layers from network inputs, and 'reverse' will start "
+            "from the network outputs",
+            choices=["forward", "reverse"],
+            default="forward",
+            dest="direction",
+        )
+        parser.add_argument(
+            "-p",
+            "--precision",
+            help="Precision to use when marking layers to run in higher precision",
+            choices=["float32", "float16"],
+            default="float32",
+        )
+
+    def setup(self, args, network):
+        self.precision = {"float32": trt.float32, "float16": trt.float16}[
+            args.precision
+        ]
+
+        if self.precision == trt.float16 and not self.arg_groups[TrtConfigArgs].fp16:
+            G_LOGGER.critical(
+                "Cannot mark layers to run in float16 if it is not enabled in the builder configuration.\n"
+                "Please also specify `--fp16` as a command-line option"
+            )
+
+        if self.precision == trt.float16 and not self.arg_groups[TrtConfigArgs].int8:
+            G_LOGGER.warning(
+                "Using float16 as the higher precision, but float16 is also the lowest precision available. "
+                "Did you mean to set --int8 as well?"
+            )
+
+        if not any(
+            [
+                self.arg_groups[TrtConfigArgs].tf32,
+                self.arg_groups[TrtConfigArgs].fp16,
+                self.arg_groups[TrtConfigArgs].int8,
+            ]
+        ):
+            G_LOGGER.critical(
+                "Please enable at least one precision besides float32 (e.g. --int8, --fp16, --tf32)"
+            )
+
+        if self.arg_groups[ModelArgs].model_type == "engine":
+            G_LOGGER.critical(
+                "The precision tool cannot work with engines, as they cannot be modified. "
+                "Please provide a different format, such as an ONNX model or TensorRT network script."
+            )
+
+        G_LOGGER.start(f"Using {self.precision} as higher precision")
+
+        if args.mode == "linear":
+            self.layer_marker = LinearMarker(len(network), args.direction)
+        elif args.mode == "bisect":
+            self.layer_marker = BisectMarker(len(network), args.direction)
+
+        self.original_precisions = {}
+        try:
+            for index, layer in enumerate(network):
+                if layer.precision_is_set:
+                    self.original_precisions[index] = layer.precision
+        except AttributeError:
+            trt_util.fail_unavailable("layer precision in polygraphy debug precision")
+
+    def mark_layers(self, network, indices):
+        EXCLUDE_LAYER_NAMES = ["CONSTANT"]
+        EXCLUDE_LAYERS = [
+            getattr(trt.LayerType, attr)
+            for attr in EXCLUDE_LAYER_NAMES
+            if hasattr(trt.LayerType, attr)
+        ]
+
+        # First, reset, since changes from the previous call will persist.
+        try:
+            for index, layer in enumerate(network):
+                if index in self.original_precisions:
+                    layer.precision = self.original_precisions[index]
+                else:
+                    layer.reset_precision()
+        except AttributeError:
+            trt_util.fail_unavailable("layer precision in polygraphy debug precision")
+
+        marked_indices = set()
+        for index in indices:
+            layer = network.get_layer(index)
+
+            def should_exclude():
+                has_non_execution_output = any(
+                    not layer.get_output(i).is_execution_tensor
+                    for i in range(layer.num_outputs)
+                )
+                has_non_activation_output = any(
+                    layer.get_output(i).dtype
+                    not in [trt.float32, trt.float16, trt.int8]
+                    for i in range(layer.num_outputs)
+                )
+                return (
+                    layer.type in EXCLUDE_LAYERS
+                    or has_non_execution_output
+                    or has_non_activation_output
+                )
+
+            if not should_exclude():
+                G_LOGGER.extra_verbose(
+                    f"Running layer in higher precision: {trt_util.str_from_layer(layer, index)}"
+                )
+                try:
+                    layer.precision = self.precision
+                except AttributeError:
+                    trt_util.fail_unavailable(
+                        "layer precision in polygraphy debug precision"
+                    )
+                marked_indices.add(index)
+
+        G_LOGGER.verbose(
+            f"Marking layer(s): {marked_indices} to run in {self.precision} precision"
+        )
+
+    def process_network(self, network):
+        indices = list(self.layer_marker.select_layers())
+        self.mark_layers(network, indices)
+
+    def step(self, success):
+        return self.layer_marker.step(success)
+
+    def remaining(self):
+        return self.layer_marker.remaining()