def __call__(self):
"""
save_uff (bool): Whether to write the generated UFF and corresponding PBTXT files.
"""
from polygraphy.backend.tf import util as tf_util
import uff
misc.log_module_info(uff)
graph, output_names = self.tf_loader()
output_names = [name.split(":")[0] for name in output_names]
# GraphDefs don't have names, so we have to name it something generic.
output_filename = None if not self.uff_path else "out.uff"
# Generate the UFF model and get information about the input_buffers/output_buffers.
uff_model, input_nodes, _ = uff.from_tensorflow(
graph.as_graph_def(),
return_graph_info=True,
quiet=(G_LOGGER.severity > G_LOGGER.VERBOSE),
debug_mode=(G_LOGGER.severity == G_LOGGER.EXTRA_VERBOSE),
text=self.uff_path,
save_preprocessed=self.uff_path,
output_filename=output_filename,
preprocessor=self.preprocessor)
input_names = [node.name for node in input_nodes]
input_shapes = [
tuple(int(dim.size) for dim in node.attr["shape"].shape.dim)
for node in input_nodes
]
return uff_model, input_names, input_shapes, output_names
def __call__(self):
"""
Converts a TensorFlow model into ONNX.
Returns:
onnx.ModelProto: The ONNX model.
"""
import tensorflow as tf
import tf2onnx
from polygraphy.backend.tf import util as tf_util
misc.log_module_info(tf2onnx)
(graph, output_names), _ = misc.try_call(self._graph)
input_names = list(tf_util.get_input_metadata(graph).keys())
if self.fold_constant:
G_LOGGER.info("Folding constants in graph using tf2onnx.tfonnx.tf_optimize")
graphdef = graph.as_graph_def()
if self.optimize:
graphdef = tf2onnx.tfonnx.tf_optimize(input_names, output_names, graph.as_graph_def(), fold_constant=self.fold_constant)
with tf.Graph().as_default() as graph, tf.compat.v1.Session(graph=graph) as sess:
tf.import_graph_def(graphdef, name="")
onnx_graph = tf2onnx.tfonnx.process_tf_graph(graph, input_names=input_names, output_names=output_names, opset=self.opset)
if self.optimize:
onnx_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
return onnx_util.check_model(onnx_graph.make_model("model"))
def __call__(self):
"""
Loads an ONNX model from a file.
Returns:
onnx.ModelProto: The ONNX model
"""
import onnx
misc.log_module_info(onnx)
G_LOGGER.verbose("Loading ONNX model: {:}".format(self.path))
return onnx_util.check_model(onnx.load(self.path))
def __call__(self, args):
misc.log_module_info(polygraphy)
script = build_script(args)
if args.gen_script:
with args.gen_script:
args.gen_script.write(script)
path = args.gen_script.name
# Somehow, piping fools isatty, e.g. `polygraphy run --gen-script - | cat`
if not args.gen_script.isatty() and path not in [
"<stdout>", "<stderr>"
]:
G_LOGGER.info("Writing script to: {:}".format(path))
# Make file executable
os.chmod(path, os.stat(path).st_mode | 0o111)
else:
exec(script)
return 0
def run(self, args):
if self.makers[TrtLoaderArgs].network_api and not tools_util.get(
args, "gen_script"):
G_LOGGER.critical(
"Cannot use the --network-api option if --gen/--gen-script is not being used."
)
elif self.makers[
TrtLoaderArgs].network_api and "trt" not in args.runners:
args.runners.append("trt")
if self.makers[
ModelArgs].model_file is None and args.runners and self.makers[
TrtLoaderArgs].network_api is None:
G_LOGGER.critical(
"One or more runners was specified, but no model file was provided. Make sure you've specified the model path, "
"and also that it's not being consumed as an argument for another parameter"
)
misc.log_module_info(polygraphy)
script = self.build_script(args)
if args.gen_script:
with args.gen_script:
args.gen_script.write(script)
path = args.gen_script.name
# Somehow, piping fools isatty, e.g. `polygraphy run --gen-script - | cat`
if not args.gen_script.isatty() and path not in [
"<stdout>", "<stderr>"
]:
G_LOGGER.info("Writing script to: {:}".format(path))
# Make file executable
os.chmod(path, os.stat(path).st_mode | 0o111)
else:
exec(script)
return 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
import tensorrt as trt
from polygraphy.backend.base import BaseRunner
from polygraphy.backend.trt import util as trt_util
from polygraphy.backend.trt.buffers import Buffers
from polygraphy.logger.logger import G_LOGGER
from polygraphy.util import cuda, misc
misc.log_module_info(trt)
class TrtRunner(BaseRunner):
"""
Runs inference using a TensorRT engine.
"""
def __init__(self, engine, name=None):
"""
Args:
engine (Callable() -> trt.ICudaEngine):
A callable that can supply a TensorRT engine.
If instead of a loader, the engine is provided directly, then the runner
will *not* take ownership of it, and therefore will not destroy it.
# 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 onnxruntime
from polygraphy.backend.base import BaseLoadModel
from polygraphy.util import misc
misc.log_module_info(onnxruntime)
class SessionFromOnnxBytes(BaseLoadModel):
def __init__(self, model_bytes):
"""
Functor that builds an ONNX-Runtime inference session.
Args:
model_bytes (Callable() -> bytes): A loader that can supply a serialized ONNX model.
"""
self._model_bytes = model_bytes
def __call__(self):
"""
Builds an ONNX-Runtime inference session.
# 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
import tensorflow as tf
from polygraphy.backend.base import BaseRunner
from polygraphy.backend.tf import util as tf_util
from polygraphy.logger.logger import G_LOGGER
from polygraphy.util import misc
misc.log_module_info(tf)
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.
timeline_dir (str):