# # 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. # from polygraphy import mod, util from polygraphy.logger import G_LOGGER, LogMode from polygraphy.tools.args import util as args_util from polygraphy.tools.args.base import BaseArgs from polygraphy.tools.args.model import ModelArgs from polygraphy.tools.script import make_invocable @mod.export() class TfTrtArgs(BaseArgs): """ [UNTESTED] TensorFlow-TensorRT Integration: TensorFlow-TensorRT. Depends on: - TrtConfigArgs """ def add_parser_args_impl(self): self.group.add_argument( "--tftrt", "--use-tftrt", help="Enable TF-TRT integration", action="store_true", default=None, dest="use_tftrt", ) self.group.add_argument( "--minimum-segment-size", help="Minimum length of a segment to convert to TensorRT", type=int, default=None, ) self.group.add_argument( "--dynamic-op", help="Enable dynamic mode (defers engine build until runtime)", action="store_true", default=None, ) def parse_impl(self, args): """ Parses command-line arguments and populates the following attributes: Attributes: use_tftrt (bool): Whether to use TF-TRT. minimum_segment_size (int): The minimum size of segments offloaded to TRT. dynamic_op (bool): Whether to enable dynamic mode, which defers engine building until runtime. """ self.use_tftrt = args_util.get(args, "use_tftrt") self.minimum_segment_size = args_util.get(args, "minimum_segment_size") self.dynamic_op = args_util.get(args, "dynamic_op") def add_to_script_impl(self, script, loader_name=None): """ Args: loader_name (str): The name of the loader which should be consumed by the ``UseTfTrt`` loader. """ if self.use_tftrt: from polygraphy.tools.args.backend.trt import TrtConfigArgs script.add_import(imports=["UseTfTrt"], frm="polygraphy.backend.tf") loader_str = make_invocable( "UseTfTrt", loader_name, max_workspace_size=self.arg_groups[TrtConfigArgs]._workspace, fp16=self.arg_groups[TrtConfigArgs].fp16, int8=self.arg_groups[TrtConfigArgs].int8, is_dynamic_op=self.dynamic_op, minimum_segment_size=self.minimum_segment_size, ) loader_name = script.add_loader(loader_str, "use_tftrt") return loader_name @mod.export() class TfLoadArgs(BaseArgs): """ TensorFlow Model Loading: loading TensorFlow models. Depends on: - ModelArgs - TfTrtArgs: if allow_tftrt == True - TrtSaveEngineBytesArgs: if allow_tftrt == True """ def __init__( self, allow_artifacts: bool = None, allow_custom_outputs: bool = None, allow_tftrt: bool = None, ): """ Args: allow_artifacts (bool): Whether to allow saving artifacts to the disk, like frozen models or TensorBoard visualizations. Defaults to True. allow_custom_outputs (bool): Whether to allow marking custom output tensors. Defaults to True. allow_tftrt (bool): Whether to allow applying TF-TRT. Defaults to False. """ super().__init__() self._allow_artifacts = util.default(allow_artifacts, True) self._allow_custom_outputs = util.default(allow_custom_outputs, True) self._allow_tftrt = util.default(allow_tftrt, False) def add_parser_args_impl(self): self.group.add_argument( "--ckpt", help="[EXPERIMENTAL] Name of the checkpoint to load. Required if the `checkpoint` file is missing. Should not include file extension " "(e.g. to load `model.meta` use `--ckpt=model`)", default=None, ) if self._allow_custom_outputs: self.group.add_argument( "--tf-outputs", help="Name(s) of TensorFlow output(s). " "Using '*' indicates that all tensors should be used as outputs.", nargs="+", default=None, ) if self._allow_artifacts: self.group.add_argument( "--save-pb", help="Path to save the TensorFlow frozen graphdef", default=None, dest="save_frozen_graph_path", ) self.group.add_argument( "--save-tensorboard", help="[EXPERIMENTAL] Path to save a TensorBoard visualization", default=None, dest="save_tensorboard_path", ) self.group.add_argument( "--freeze-graph", help="[EXPERIMENTAL] Attempt to freeze the graph", action="store_true", default=None, ) def parse_impl(self, args): """ Parses command-line arguments and populates the following attributes: Attributes: ckpt (str): Name of the checkpoint. outputs (List[str]): Names of output tensors. save_frozen_graph_path (str): The path at which the frozen graph will be saved. save_tensorboard_path (str): The path at which the TensorBoard visualization will be saved. freeze_graph (bool): Whether to attempt to freeze the graph. """ self.ckpt = args_util.get(args, "ckpt") self.outputs = args_util.get_outputs(args, "tf_outputs") self.save_frozen_graph_path = args_util.get(args, "save_frozen_graph_path") self.save_tensorboard_path = args_util.get(args, "save_tensorboard_path") self.freeze_graph = args_util.get(args, "freeze_graph") def add_to_script_impl(self, script, disable_custom_outputs=None): """ Args: disable_custom_outputs (bool): Whether to disallow modifying outputs according to the `outputs` attribute. Defaults to False. """ model_file = self.arg_groups[ModelArgs].path model_type = self.arg_groups[ModelArgs].model_type if model_type == "ckpt": G_LOGGER.verbose( f"Loading a TensorFlow checkpoint from {model_file}. Please ensure you are not using the --use-subprocess flag", mode=LogMode.ONCE, ) script.add_import(imports=["GraphFromCkpt"], frm="polygraphy.backend.tf") loader_id = "load_ckpt" loader_str = make_invocable("GraphFromCkpt", model_file, self.ckpt) elif model_type == "keras": script.add_import(imports=["GraphFromKeras"], frm="polygraphy.backend.tf") loader_id = "load_keras" loader_str = make_invocable("GraphFromKeras", model_file) elif model_type == "frozen": script.add_import(imports=["GraphFromFrozen"], frm="polygraphy.backend.tf") G_LOGGER.verbose( "Attempting to load as a frozen graph. If this is not correct, please specify --model-type", mode=LogMode.ONCE, ) loader_id = "load_frozen" loader_str = make_invocable("GraphFromFrozen", model_file) else: G_LOGGER.critical( f"Model type: {model_type} cannot be imported with TensorFlow." ) loader_name = script.add_loader(loader_str, loader_id) if self.freeze_graph: script.add_import(imports=["OptimizeGraph"], frm="polygraphy.backend.tf") loader_name = script.add_loader( make_invocable("OptimizeGraph", loader_name), "optimize_graph" ) engine_dir = None if self._allow_tftrt: from polygraphy.tools.args.backend.trt import TrtSaveEngineBytesArgs loader_name = self.arg_groups[TfTrtArgs].add_to_script(script, loader_name) engine_dir = self.arg_groups[TrtSaveEngineBytesArgs].path MODIFY_TF = "ModifyGraphOutputs" outputs = ( None if disable_custom_outputs else args_util.get_outputs_for_script(script, self.outputs) ) modify_tf_str = make_invocable(MODIFY_TF, loader_name, outputs=outputs) if modify_tf_str != make_invocable(MODIFY_TF, loader_name): script.add_import(imports=[MODIFY_TF], frm="polygraphy.backend.tf") loader_name = script.add_loader(modify_tf_str, "modify_tf") WRITE_TF = "SaveGraph" write_tf_str = make_invocable( WRITE_TF, loader_name, path=self.save_frozen_graph_path, tensorboard_dir=self.save_tensorboard_path, engine_dir=engine_dir, ) if write_tf_str != make_invocable(WRITE_TF, loader_name): script.add_import(imports=[WRITE_TF], frm="polygraphy.backend.tf") loader_name = script.add_loader(write_tf_str, "save_tf") return loader_name def load_graph(self): """ Loads a TensorFlow graph according to arguments provided on the command-line. Returns: tf.Graph """ loader = args_util.run_script(self.add_to_script) return loader()