def set_trt_logging_level(sev):
from polygraphy import mod
trt = mod.lazy_import("tensorrt")
if not mod.has_mod(trt, with_attr="__version__"):
return
if sev >= G_LOGGER.CRITICAL:
get_trt_logger().min_severity = trt.Logger.INTERNAL_ERROR
elif sev >= G_LOGGER.ERROR:
get_trt_logger().min_severity = trt.Logger.ERROR
elif sev >= G_LOGGER.INFO:
get_trt_logger().min_severity = trt.Logger.WARNING
elif sev >= G_LOGGER.VERBOSE:
get_trt_logger().min_severity = trt.Logger.INFO
else:
get_trt_logger().min_severity = trt.Logger.VERBOSE
def set_tf_logging_level(sev):
import os
from polygraphy import mod
tf = mod.lazy_import("tensorflow", version="<2.0")
if not mod.has_mod(tf, with_attr="__version__"):
return
if sev > G_LOGGER.WARNING:
tf_sev = tf.compat.v1.logging.ERROR
tf_logging_level = "3"
elif sev > G_LOGGER.INFO:
tf_sev = tf.compat.v1.logging.WARN
tf_logging_level = "2"
elif sev > G_LOGGER.VERBOSE:
tf_sev = tf.compat.v1.logging.INFO
tf_logging_level = "1"
else:
tf_sev = tf.compat.v1.logging.DEBUG
tf_logging_level = "0"
tf.compat.v1.logging.set_verbosity(tf_sev)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = tf_logging_level
# 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
from polygraphy.logger import G_LOGGER
from polygraphy.tools.base import Tool
algorithm_selector = mod.lazy_import("polygraphy.backend.trt.algorithm_selector")
class Tactics(Tool):
"""
Display the contents of tactic replay files in a human readable format.
(for example, those generated by `--save-tactics` from `polygraphy run`)
"""
def __init__(self):
super().__init__("tactics")
def add_parser_args(self, parser):
parser.add_argument("tactic_replay", help="Path to a tactic replay file")
def run(self, args):
# 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 base64
import functools
import io
import json
from collections import OrderedDict
from polygraphy import config, constants, mod
from polygraphy.logger import G_LOGGER
np = mod.lazy_import("numpy")
util = mod.lazy_import("polygraphy.util.util")
TYPE_STRING_PREFIX = "__polygraphy_encoded_"
def legacy_str_from_type(typ):
return TYPE_STRING_PREFIX + typ.__name__
def str_from_type(typ):
return typ.__name__
class BaseCustomImpl(object):
"""
# 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 copy
from polygraphy import mod
from polygraphy.common import TensorMetadata
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import DataLoaderArgs, ModelArgs, OnnxLoaderArgs, OnnxSaveArgs, OnnxShapeInferenceArgs
from polygraphy.tools.args import util as args_util
from polygraphy.tools.surgeon.subtool.base import BaseSurgeonSubtool
onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
class Extract(BaseSurgeonSubtool):
"""
Extract a subgraph from an ONNX model based on the specified inputs and outputs.
"""
def __init__(self):
super().__init__("extract")
self.subscribe_args(
ModelArgs(
model_required=True,
inputs="--model-inputs",
model_type="onnx",
# 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
from polygraphy.common.interface import TypedDict
np = mod.lazy_import("numpy")
class MetadataTuple(object):
def __init__(self, dtype, shape):
self.dtype = dtype
self.shape = shape
def __iter__(self):
yield from [self.dtype, self.shape]
def __repr__(self):
return "MetadataTuple({:}, {:})".format(self.dtype, self.shape)
def __str__(self):
ret = ""
#
# 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
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import (DataLoaderArgs, ModelArgs, OnnxLoaderArgs,
OnnxSaveArgs, OnnxShapeInferenceArgs)
from polygraphy.tools.args import util as args_util
from polygraphy.tools.args.base import BaseArgs
from polygraphy.tools.surgeon.subtool.base import BaseSurgeonSubtool
gs = mod.lazy_import("onnx_graphsurgeon")
class OnnxNodeArgs(BaseArgs):
def add_to_parser(self, parser):
node_args = parser.add_argument_group(
"Inserted Node", "Options for the node to insert")
node_args.add_argument(
"--inputs",
help=
"The names of input tensors for the new node. Order will be preserved. "
"Format: --inputs <name>. For example: --inputs name0 name1",
nargs="+",
required=True)
node_args.add_argument(
"--outputs",
# 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
from polygraphy.tools import util as tools_util
from polygraphy.tools.args import DataLoaderArgs, ModelArgs, OnnxLoaderArgs, OnnxSaveArgs, OnnxShapeInferenceArgs
from polygraphy.tools.surgeon.subtool.base import BaseSurgeonSubtool
onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
gs = mod.lazy_import("onnx_graphsurgeon")
class Sanitize(BaseSurgeonSubtool):
"""
Clean up, optimize, and/or change input shapes in an ONNX model.
"""
def __init__(self):
super().__init__("sanitize")
self.subscribe_args(
ModelArgs(
model_required=True,
inputs="--override-inputs",
# 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 constants, mod, util
from polygraphy.common import TensorMetadata
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import util as args_util
from polygraphy.tools.args.base import BaseArgs
from polygraphy.tools.script import Script, make_invocable
onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
onnxrt_backend = mod.lazy_import("polygraphy.backend.onnxrt")
@mod.export()
class OnnxSaveArgs(BaseArgs):
def __init__(self,
infer_shapes=False,
output="output",
short_opt="-o",
required=False):
super().__init__()
self._infer_shapes = infer_shapes
self._output = output
self._short_opt = short_opt
self._required = required
#
# 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 os
from polygraphy import mod
from polygraphy.common.interface import TypedDict
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import ModelArgs, OnnxLoaderArgs, OnnxSaveArgs, OnnxShapeInferenceArgs
from polygraphy.tools.base import Tool
common_backend = mod.lazy_import("polygraphy.backend.common")
gs = mod.lazy_import("onnx_graphsurgeon")
onnx_backend = mod.lazy_import("polygraphy.backend.onnx")
onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
trt = mod.lazy_import("tensorrt")
trt_backend = mod.lazy_import("polygraphy.backend.trt")
trt_util = mod.lazy_import("polygraphy.backend.trt.util")
util = mod.lazy_import("polygraphy.util")
class UnsupportedNodeDict(TypedDict(lambda: str, lambda: dict)):
"""
An ordered dictionary that maps ops to error(s) encountered by TensorRT
while trying to parse them, and the range of node indices for the subgraphs
where these errors were encountered.
#
# 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.logger import G_LOGGER
from polygraphy.tools.args import ModelArgs, TrtConfigArgs
from polygraphy.tools.debug.subtool.base import BaseCheckerSubtool
trt = mod.lazy_import("tensorrt")
trt_util = mod.lazy_import("polygraphy.backend.trt.util")
class BaseMarker(object):
def __init__(self, max_layers, direction, num_layers):
self.max_layers = max_layers
self.direction = direction
self.num_layers = num_layers
self.good = max_layers + 1 # Pretend marking all the layers gives us good accuracy.
self.iteration = 0
def select_layers(self):
self.iteration += 1
if self.direction == "forward":
G_LOGGER.info("Selecting first {:} layer(s) to run in higher precision".format(self.num_layers))
# 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 copy
import time
from collections import OrderedDict
from polygraphy import mod, util
from polygraphy.backend.base import BaseRunner
from polygraphy.backend.pluginref.references import OP_REGISTRY
from polygraphy.logger import G_LOGGER
np = mod.lazy_import("numpy")
onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
@mod.export()
class PluginRefRunner(BaseRunner):
"""
Runner for reference checking TensorRT plugins
"""
def __init__(self, graph, name=None):
"""
Args:
graph (Union[onnx_graphsurgeon.Graph, Callable() -> onnx_graphsurgeon.Graph]):
An ONNX-GraphSurgeon graph or a callable that returns one.
name (str):
The human-readable name prefix to use for this runner.
# 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 copy
from collections import OrderedDict
from polygraphy.logger.logger import G_LOGGER
from polygraphy import mod
util = mod.lazy_import("polygraphy.util")
def TypedDict(key_type_func, value_type_func):
"""
Returns a class (not an instance) that will provide a dictionary-like
interface with runtime type checks.
Note: The types are provided lazily via a callable to avoid unnecessary dependencies
on types from external packages at import-time.
Args:
key_type_func (Callable() -> type):
A callable that returns the expected key type.
value_type_func (Callable() -> type):
A callable that returns the expected value type.
# 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
import time
from collections import OrderedDict
from polygraphy import func, mod, util
from polygraphy.backend.base import BaseRunner
from polygraphy.backend.tf import util as tf_util
from polygraphy.logger import G_LOGGER
tf = mod.lazy_import("tensorflow", version="<2.0")
@mod.export()
class TfRunner(BaseRunner):
"""
Runs inference using a TensorFlow session.
"""
def __init__(self, sess, timeline_dir=None, name=None):
"""
Args:
sess (Callable() -> Tuple[tf.Session, Sequence[str]]):
A callable that can supply a tuple containing a
TensorFlow session and output names.
import contextlib
from polygraphy import mod, util
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import (
ModelArgs,
OnnxLoaderArgs,
OnnxShapeInferenceArgs,
TfLoaderArgs,
TrtEngineLoaderArgs,
TrtNetworkLoaderArgs,
TrtPluginLoaderArgs,
)
from polygraphy.tools.base import Tool
trt_util = mod.lazy_import("polygraphy.backend.trt.util")
onnx_util = mod.lazy_import("polygraphy.backend.onnx.util")
tf_util = mod.lazy_import("polygraphy.backend.tf.util")
class Model(Tool):
"""
Display information about a model, including inputs and outputs, as well as layers and their attributes.
"""
def __init__(self):
super().__init__("model")
self.subscribe_args(ModelArgs(model_required=True, inputs=None))
self.subscribe_args(TfLoaderArgs(artifacts=False, outputs=False))
self.subscribe_args(OnnxShapeInferenceArgs())
self.subscribe_args(OnnxLoaderArgs(output_prefix=None))
self.subscribe_args(TrtPluginLoaderArgs())
# 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 copy
import os
import sys
import tempfile
from polygraphy import constants, mod, util
from polygraphy.backend.base import BaseLoader
from polygraphy.backend.onnx import util as onnx_util
from polygraphy.logger import G_LOGGER, LogMode
onnx = mod.lazy_import("onnx", version=">=1.8.1")
onnxrt = mod.lazy_import("onnxruntime")
onnxmltools = mod.lazy_import("onnxmltools")
tf = mod.lazy_import("tensorflow", version="<2.0")
tf2onnx = mod.lazy_import("tf2onnx")
tf_util = mod.lazy_import("polygraphy.backend.tf.util", log=False)
gs = mod.lazy_import("onnx_graphsurgeon", version=mod.LATEST_VERSION)
shape_inference = mod.lazy_import("onnx.shape_inference")
external_data_helper = mod.lazy_import("onnx.external_data_helper")
LARGE_MODEL_THRESHOLD = 512 << 20 # 512 MiB
class BaseLoadOnnxCopy(BaseLoader):
"""
Abstract base class for loaders that require loading an ONNX model and potentially
# 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 copy
from collections import OrderedDict
from polygraphy import mod, util
from polygraphy.common import TensorMetadata
from polygraphy.logger import G_LOGGER
onnx = mod.lazy_import("onnx")
numpy_helper = mod.lazy_import("onnx.numpy_helper")
def get_num_nodes(model):
def _get_num_graph_nodes(graph):
num_nodes = len(graph.node)
for node in graph.node:
for attr in node.attribute:
if attr.type == onnx.AttributeProto.GRAPH:
num_nodes += _get_num_graph_nodes(attr.g)
elif attr.type == onnx.AttributeProto.GRAPHS:
for subgraph in attr.graphs:
num_nodes += _get_num_graph_nodes(subgraph)
return num_nodes
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import (
DataLoaderArgs,
ModelArgs,
OnnxLoaderArgs,
OnnxShapeInferenceArgs,
TrtConfigArgs,
TrtEngineLoaderArgs,
TrtEngineSaveArgs,
TrtNetworkLoaderArgs,
TrtPluginLoaderArgs,
)
from polygraphy.tools.base import Tool
from polygraphy.tools.debug.subtool.artifact_sorter import ArtifactSorterArgs
trt_backend = mod.lazy_import("polygraphy.backend.trt")
class BaseCheckerSubtool(Tool):
def __init__(self, name, strict_types_default=None, prefer_artifacts=True):
super().__init__(name)
self.subscribe_args(
ArtifactSorterArgs("polygraphy_debug.engine",
prefer_artifacts=prefer_artifacts))
self.subscribe_args(ModelArgs(model_required=True, inputs=None))
self.subscribe_args(OnnxShapeInferenceArgs())
self.subscribe_args(OnnxLoaderArgs(output_prefix=None))
self.subscribe_args(DataLoaderArgs()) # For int8 calibration
self.subscribe_args(
TrtConfigArgs(strict_types_default=strict_types_default))
self.subscribe_args(TrtPluginLoaderArgs())
#
# 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 time
from collections import OrderedDict
from polygraphy import func, mod, util
from polygraphy.backend.base import BaseRunner
torch = mod.lazy_import("torch")
@mod.export()
class PytRunner(BaseRunner):
"""
Runs inference using PyTorch.
"""
def __init__(self, model, input_metadata, output_names, name=None):
"""
Args:
model (Callable() -> torch.nn.Module):
A model loader that returns a torch.nn.Module or subclass.
input_metadata (TensorMetadata): Mapping of input names to their data types and shapes.
output_names (List[str]):
# 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 copy
from collections import OrderedDict
from polygraphy import mod, util
from polygraphy.common import TensorMetadata
from polygraphy.logger import G_LOGGER
gs = mod.lazy_import("onnx_graphsurgeon")
numpy_helper = mod.lazy_import("onnx.numpy_helper")
onnx = mod.lazy_import("onnx")
def get_num_nodes(model):
def _get_num_graph_nodes(graph):
num_nodes = len(graph.node)
for node in graph.node:
for attr in node.attribute:
if attr.type == onnx.AttributeProto.GRAPH:
num_nodes += _get_num_graph_nodes(attr.g)
elif attr.type == onnx.AttributeProto.GRAPHS:
for subgraph in attr.graphs:
num_nodes += _get_num_graph_nodes(subgraph)
return num_nodes
# 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 contextlib
import copy
import ctypes
import time
from polygraphy import constants, mod, util
from polygraphy.backend.base import BaseLoader
from polygraphy.backend.trt import util as trt_util
from polygraphy.backend.trt.profile import Profile
from polygraphy.logger import G_LOGGER
trt = mod.lazy_import("tensorrt")
gs = mod.lazy_import("onnx_graphsurgeon")
np = mod.lazy_import("numpy")
@mod.export(funcify=True)
class LoadPlugins(BaseLoader):
"""
A passthrough loader that loads plugins from the specified paths.
Passthrough here means that it can be used to wrap any other loader. The purpose of wrapping
another loader is that you can control the order of execution when lazily evaluating.
For immediate evaluation, use `load_plugins` instead:
::
load_plugins(plugins=["/path/to/my/plugin.so", "/path/to/my/other_plugin.so"])
# 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 func, mod, util
from polygraphy.backend.trt import util as trt_util
from polygraphy.common.interface import TypedDict
from polygraphy.json import Decoder, Encoder, add_json_methods
from polygraphy.logger import G_LOGGER
trt = mod.lazy_import("tensorrt")
##
## Data Structures
##
# NOTE: Modifying the structure of the data classes below will break backwards compatiblity
@mod.export()
class Algorithm(object):
"""
Represents a TensorRT algorithm variant, which can be uniquely represented
by an implementation ID and tactic ID.
"""
@staticmethod
# 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.backend.base import BaseLoader
onnxruntime = mod.lazy_import("onnxruntime")
@mod.export_deprecated_alias("SessionFromOnnxBytes", remove_in="0.32.0")
@mod.export(funcify=True)
class SessionFromOnnx(BaseLoader):
"""
Functor that builds an ONNX-Runtime inference session.
"""
def __init__(self, model_bytes):
"""
Builds an ONNX-Runtime inference session.
Args:
model_bytes (Callable() -> Union[bytes, str]):
A serialized ONNX model or a path to a model, or a callable that returns the same.
#
# 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 contextlib
from collections import OrderedDict
from polygraphy import cuda, mod, util
from polygraphy.logger import G_LOGGER, LogMode
trt = mod.lazy_import("tensorrt")
np = mod.lazy_import("numpy")
@mod.export()
def Calibrator(data_loader,
cache=None,
BaseClass=None,
batch_size=None,
quantile=None,
regression_cutoff=None,
algo=None):
"""
Supplies calibration data to TensorRT to calibrate the network for INT8 inference.
Args:
# 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
from polygraphy.logger import G_LOGGER
np = mod.lazy_import("numpy")
gs = mod.lazy_import("onnx_graphsurgeon")
OP_REGISTRY = {
} # Dict[str, Callable]: Maps op names to reference implementations
def register(op):
"""
Registers a function as the reference implementation for a given op.
Args:
op (str): The name of the op for which to register this function.
"""
def register_impl(func):
def wrapped_func(node, intermediate_tensors):
