chore: import upstream snapshot with attribution

tools/Polygraphy/polygraphy/backend/tf/loader.py

@@ -0,0 +1,496 @@
+#
+# 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.
+#
+# Sets up everything needed to perform inference in TensorFlow.
+import os
+
+from polygraphy import constants, mod, util
+from polygraphy.backend.base import BaseLoader
+from polygraphy.backend.tf import util as tf_util
+from polygraphy.logger import G_LOGGER
+
+tf = mod.lazy_import("tensorflow<2.0")
+
+
+@mod.export(funcify=True)
+class OptimizeGraph(BaseLoader):
+    """
+    Functor that freezes a TensorFlow graph, and folds constants.
+    """
+
+    def __init__(self, graph):
+        """
+        Freezes a TensorFlow graph and folds constants.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+        """
+        self._graph = graph
+
+    def constfold(self, graphdef, output_names):
+        from tensorflow.core.protobuf import (
+            config_pb2,
+            meta_graph_pb2,
+            rewriter_config_pb2,
+        )
+        from tensorflow.python.framework import importer, ops
+        from tensorflow.python.grappler import tf_optimizer
+        from tensorflow.python.training import saver
+
+        graph = ops.Graph()
+        with graph.as_default():
+            output_collection = meta_graph_pb2.CollectionDef()
+            output_list = output_collection.node_list.value
+            for output in output_names:
+                output_list.append(output.encode("utf-8"))
+
+            importer.import_graph_def(graphdef, name="")
+            metagraph = saver.export_meta_graph(
+                graph_def=graph.as_graph_def(add_shapes=True), graph=graph
+            )
+            metagraph.collection_def["train_op"].CopyFrom(output_collection)
+
+        rewriter_config = rewriter_config_pb2.RewriterConfig()
+        rewriter_config.optimizers.extend(["constfold"])
+        rewriter_config.meta_optimizer_iterations = (
+            rewriter_config_pb2.RewriterConfig.ONE
+        )
+
+        session_config = config_pb2.ConfigProto()
+        session_config.graph_options.resave_options.CopyFrom(rewriter_config)
+        return tf_optimizer.OptimizeGraph(session_config, metagraph, graph_id=b"graph")
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, output_names), _ = util.invoke_if_callable(self._graph)
+        with tf.Session(graph=graph) as sess:
+            sess.run(tf.initializers.global_variables())
+            sess.run(tf.initializers.local_variables())
+
+            graphdef = sess.graph.as_graph_def()
+            removed = tf.graph_util.remove_training_nodes(graphdef)
+            G_LOGGER.ultra_verbose(f"Removed nodes: {removed}")
+
+            for node in graphdef.node:
+                if node.op == "RefSwitch":
+                    node.op = "Switch"
+                    for index in range(len(node.input)):
+                        if "moving_" in node.input[index]:
+                            node.input[index] = node.input[index] + "/read"
+                elif node.op == "AssignSub":
+                    node.op = "Sub"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                elif node.op == "AssignAdd":
+                    node.op = "Add"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                elif node.op == "Assign":
+                    node.op = "Identity"
+                    if "use_locking" in node.attr:
+                        del node.attr["use_locking"]
+                    if "validate_shape" in node.attr:
+                        del node.attr["validate_shape"]
+                    if len(node.input) == 2:
+                        # input0: ref: Should be from a Variable node. May be uninitialized.
+                        # input1: value: The value to be assigned to the variable.
+                        node.input[0] = node.input[1]
+                        del node.input[1]
+
+            # Strip port information from outputs
+            output_names = [name.split(":")[0] for name in output_names]
+            output_graph_def = tf.graph_util.convert_variables_to_constants(
+                sess, graphdef, output_names
+            )
+            output_graph_def = self.constfold(output_graph_def, output_names)
+            return graph_from_frozen(output_graph_def)
+
+
+@mod.export(funcify=True)
+class GraphFromKeras(BaseLoader):
+    """
+    Functor that loads a TensorFlow model from Keras.
+    """
+
+    def __init__(self, path):
+        """
+        Loads a TensorFlow model from Keras.
+
+        Args:
+            path (Union[str, h5py.File]): A path to the saved model, or the file object.
+        """
+        self.path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+
+        from tensorflow.python import keras
+        from tensorflow.python.keras import backend
+
+        model = keras.models.load_model(self.path)
+        graph = backend.get_session().graph
+        return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class GraphFromFrozen(BaseLoader):
+    """
+    Functor that loads a TensorFlow frozen model.
+    """
+
+    def __init__(self, path):
+        """
+        Loads a TensorFlow frozen model.
+
+        Args:
+            path (Union[str, tf.Graph, tf.GraphDef]):
+                    A path to the frozen model, or a frozen TensorFlow graph or graphdef.
+        """
+        self.path = path
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        graph = tf_util.load_graph(self.path)
+        return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class GraphFromCkpt(BaseLoader):
+    """
+    Functor that loads a TensorFlow model from a checkpoint. Note that in order to use checkpoints,
+    you must NOT use subprocesses in the Comparator.
+    """
+
+    def __init__(self, dir, name=None):
+        """
+        Loads a TensorFlow model from a checkpoint.
+
+        Args:
+            dir (str): Path to a directory containing checkpoints.
+
+
+            name (str):
+                    The name of the checkpoint to load, not including the file extension.
+                    For example, to load `model.meta`, the argument would be `model`.
+        """
+        self.dir = dir
+        self.name = name
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        # If `name` is not provided, this expects that the directory contains a `checkpoint` file with the contents:
+        #
+        # model_checkpoint_path: "model"
+        # all_model_checkpoint_paths: "model"
+        #
+        # where "model" is the checkpoint name
+        if not os.path.isdir(self.dir):
+            G_LOGGER.warning(
+                f"Specified checkpoint directory: {self.dir} does not look like a directory."
+            )
+
+        if self.name is None:
+            G_LOGGER.verbose(
+                "Checkpoint name was not explicitly provided, searching for `checkpoint` file"
+            )
+            checkpoint = tf.train.get_checkpoint_state(self.dir)
+            if checkpoint is None:
+                ckpt_file_contents = '\nmodel_checkpoint_path: "model"\nall_model_checkpoint_paths: "model"\n'
+                G_LOGGER.critical(
+                    f"Checkpoint directory: {self.dir} does not contain a `checkpoint` file, and the checkpoint name was not provided. Please either create a checkpoint file with the contents:\n{ckpt_file_contents} \nWhere `model` is the name of the checkpoint, or explicitly provide the name with --ckpt, not including file extensions"
+                )
+            input_checkpoint = checkpoint.model_checkpoint_path
+        else:
+            input_checkpoint = os.path.join(self.dir, self.name)
+
+        meta_file = input_checkpoint + ".meta"
+        with tf.Graph().as_default() as graph, tf.compat.v1.Session(
+            graph=graph
+        ).as_default() as sess:
+            saver = tf.compat.v1.train.import_meta_graph(meta_file, clear_devices=True)
+            saver.restore(sess, input_checkpoint)
+            return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class UseTfTrt(BaseLoader):
+    """
+    [UNTESTED] Functor that optimizes a TensorFlow model using TF-TRT.
+    """
+
+    def __init__(
+        self,
+        graph,
+        max_workspace_size=None,
+        fp16=None,
+        int8=None,
+        max_batch_size=None,
+        is_dynamic_op=False,
+        minimum_segment_size=None,
+    ):
+        """
+        Optimizes a TensorFlow model using TF-TRT.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+            max_workspace_size (int): The maximum workspace size.
+            fp16 (bool): Whether to run in FP16 mode.
+            max_batch_size (int): The maximum batch size.
+        """
+        self._graph = graph
+        self.max_workspace_size = util.default(max_workspace_size, 1 << 24)
+        self.fp16 = util.default(fp16, False)
+        self.fp8 = util.default(fp8, False)
+        self.int8 = util.default(int8, False)
+        self.max_batch_size = util.default(max_batch_size, 1)
+        self.is_dynamic_op = is_dynamic_op
+        self.minimum_segment_size = util.default(minimum_segment_size, 3)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        from tensorflow.contrib import tensorrt as tf_trt
+
+        (graph, output_names), _ = util.invoke_if_callable(self._graph)
+
+        precision_mode = "FP16" if self.fp16 else "FP32"
+        precision_mode = "INT8" if self.int8 else precision_mode
+        precision_mode = "FP8" if self.fp8 else precision_mode
+
+        G_LOGGER.info(
+            f"For TF-TRT, using outputs={output_names}, max_workspace_size_bytes={self.max_workspace_size}, max_batch_size={self.max_batch_size}, minimum_segment_size={self.minimum_segment_size}, is_dynamic_op={self.is_dynamic_op}, precision_mode={precision_mode}"
+        )
+
+        graphdef = tf_trt.create_inference_graph(
+            graph.as_graph_def(),
+            outputs=output_names,
+            max_workspace_size_bytes=self.max_workspace_size,
+            max_batch_size=self.max_batch_size,
+            minimum_segment_size=self.minimum_segment_size,
+            is_dynamic_op=self.is_dynamic_op,
+            precision_mode=precision_mode,
+        )
+
+        segment_number = 0
+        for node in graphdef.node:
+            if node.op == "TRTEngineOp":
+                engine = node.attr["serialized_segment"].s
+                segment_number += 1
+        G_LOGGER.info(f"Found {segment_number} engines in TFTRT graph")
+
+        with tf.Graph().as_default() as graph:
+            tf.import_graph_def(graphdef, name="")
+            return graph, tf_util.get_graph_output_names(graph)
+
+
+@mod.export(funcify=True)
+class ModifyGraphOutputs(BaseLoader):
+    """
+    Functor that modifies outputs of a TensorFlow graph.
+    """
+
+    def __init__(self, graph, outputs=None):
+        """
+        Modifies outputs of a TensorFlow graph.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            outputs (List[str]):
+                    Names of output tensors. If provided, this will override the outputs
+                    determined by the loader.
+                    If a value of `constants.MARK_ALL` is used instead of a list, all tensors in the network are marked.
+        """
+        self._graph = graph
+        self.outputs = outputs
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, outputs), _ = util.invoke_if_callable(self._graph)
+
+        if self.outputs == constants.MARK_ALL:
+            outputs = list(tf_util.get_output_metadata(graph, layerwise=True).keys())
+        elif self.outputs is not None:
+            outputs = self.outputs
+
+        return graph, outputs
+
+
+@mod.export(funcify=True)
+class SaveGraph(BaseLoader):
+    """
+    Functor that writes out artifacts from a TensorFlow graph.
+    """
+
+    def __init__(self, graph, path=None, tensorboard_dir=None, engine_dir=None):
+        """
+        Writes out artifacts from a TensorFlow Graph.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            path (str): Path at which to save the frozen graphdef.
+            tensorboard_dir (str): The directory in which to write TensorBoard visualizations.
+            engine_dir (str): The directory in which to save TF-TRT engines,
+        """
+        self._graph = graph
+        self.path = path
+        self.tensorboard_dir = tensorboard_dir
+        self.engine_dir = engine_dir
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            Tuple[tf.Graph, Sequence[str]]: The TensorFlow graph, and the names of its outputs.
+        """
+        (graph, outputs), _ = util.invoke_if_callable(self._graph)
+
+        if self.path:
+            util.save_file(graph.as_graph_def().SerializeToString(), dest=self.path)
+        if self.tensorboard_dir:
+            G_LOGGER.info(f"Writing tensorboard events to {self.tensorboard_dir}")
+            train_writer = tf.compat.v1.summary.FileWriter(self.tensorboard_dir)
+            train_writer.add_graph(graph)
+
+        if self.engine_dir is not None:
+            graphdef = graph.as_graph_def()
+            segment_number = 0
+            for node in graphdef.node:
+                if node.op == "TRTEngineOp":
+                    engine = node.attr["serialized_segment"].s
+                    if self.engine_dir is not None:
+                        util.save_file(
+                            contents=engine,
+                            dest=os.path.join(
+                                self.engine_dir, f"segment-{segment_number}"
+                            ),
+                        )
+                    segment_number += 1
+
+        return graph, outputs
+
+
+@mod.export(funcify=True)
+class CreateConfig(BaseLoader):
+    """
+    Functor that creates a TensorFlow config.
+    """
+
+    def __init__(self, gpu_memory_fraction=None, allow_growth=None, use_xla=None):
+        """
+        Creates a TensorFlow config.
+
+        Args:
+            gpu_memory_fraction (float):
+                The fraction of GPU memory that will be made available to TensorFlow.
+                This should be a value between 0.0 and 1.0.
+            allow_growth (bool): Whether to allow GPU memory allocated by TensorFlow to grow.
+            use_xla (bool): Whether to attempt to enable XLA.
+        """
+        self.gpu_memory_fraction = util.default(gpu_memory_fraction, 0.9)
+        self.allow_growth = util.default(allow_growth, False)
+        self.use_xla = util.default(use_xla, False)
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            tf.ConfigProto: The TensorFlow config.
+        """
+
+        # Session configuration
+        gpu_options = tf.compat.v1.GPUOptions(
+            per_process_gpu_memory_fraction=self.gpu_memory_fraction,
+            allow_growth=self.allow_growth,
+        )
+        config = tf.compat.v1.ConfigProto(gpu_options=gpu_options)
+        if self.use_xla:
+            config.graph_options.optimizer_options.global_jit_level = (
+                tf.OptimizerOptions.ON_1
+            )
+        G_LOGGER.verbose(
+            f"Using gpu memory fraction: {self.gpu_memory_fraction}, XLA: {self.use_xla}"
+        )
+        return config
+
+
+@mod.export(funcify=True)
+class SessionFromGraph(BaseLoader):
+    """
+    Functor that creates a TensorFlow session that can be used for inference.
+    """
+
+    def __init__(self, graph, config=None):
+        """
+        Creates a TensorFlow session.
+
+        Args:
+            graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]):
+                    A tuple containing a TensorFlow graph and output names or a callable that returns one.
+
+
+            config (Union[tf.ConfigProto, Callable() -> tf.ConfigProto]):
+                    A TensorFlow ConfigProto or a callable that returns one.
+        """
+        self.graph = graph
+        self.config = util.default(config, CreateConfig())
+
+    @util.check_called_by("__call__")
+    def call_impl(self):
+        """
+        Returns:
+            tf.Session: The TensorFlow session.
+        """
+        config, _ = util.invoke_if_callable(self.config)
+        (graph, output_names), _ = util.invoke_if_callable(self.graph)
+
+        with graph.as_default() as graph, tf.compat.v1.Session(
+            graph=graph, config=config
+        ).as_default() as sess:
+            G_LOGGER.verbose(f"Using TensorFlow outputs: {output_names}")
+            G_LOGGER.extra_verbose("Initializing variables in TensorFlow Graph")
+            sess.run(tf.compat.v1.initializers.global_variables())
+            return sess, output_names