Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor Changes to configurability - Science.io Integration #848

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Minor Changes to configurability - Science.io Integration #848

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6c7d038
Update Backend to take TRTBuilder Flags
dataartist-og Jun 7, 2022
2f58cce
Add options for large pb files, serialize to disc
dataartist-og Jun 7, 2022
3dcc149
Add onnx model even with external_data_type
dataartist-og Jun 8, 2022
c4e2a1e
Fix model.graph call on string
dataartist-og Jun 8, 2022
894f8b1
Make engine module accessible
dataartist-og Jun 10, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions onnx_tensorrt/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,5 +3,6 @@
from __future__ import absolute_import

from . import backend
from . import tensorrt_engine as engine

__version__ = "8.2.1"
38 changes: 29 additions & 9 deletions onnx_tensorrt/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,8 @@ def count_trailing_ones(vals):

class TensorRTBackendRep(BackendRep):
def __init__(self, model, device,
max_workspace_size=None, serialize_engine=False, verbose=False, **kwargs):
max_workspace_size=None, serialize_engine=False, verbose=False, \
fp16=False, flags=None, external_data_format=False, engine_path=None, **kwargs):
if not isinstance(device, Device):
device = Device(device)
self._set_device(device)
Expand All @@ -45,21 +46,36 @@ def __init__(self, model, device,
self.serialize_engine = serialize_engine
self.verbose = verbose
self.dynamic = False

self.fp16 = fp16
self.builder_flags = flags
self.engine_path = engine_path

if self.fp16 and self.builder.platform_has_fast_fp16:
self.config.set_flag(trt.BuilderFlag.FP16)

if flags is not None:
for flag in flags:
self.config.set_flag(flag)

if self.verbose:
print(f'\nRunning {model.graph.name}...')
TRT_LOGGER.min_severity = trt.Logger.VERBOSE

if not isinstance(model, six.string_types):
model_str = model.SerializeToString()
else:

if external_data_format:
parser_func = self.parser.parse_from_file
model_str = model

model = onnx.load(model)
else:
if not isinstance(model, six.string_types):
model_str = model.SerializeToString()
else:
model_str = model
parser_func = self.parser.parse

if not trt.init_libnvinfer_plugins(TRT_LOGGER, ""):
msg = "Failed to initialize TensorRT's plugin library."
raise RuntimeError(msg)

if not self.parser.parse(model_str):
if not parser_func(model_str):
error = self.parser.get_error(0)
msg = "While parsing node number %i:\n" % error.node()
msg += ("%s:%i In function %s:\n[%i] %s" %
Expand Down Expand Up @@ -140,6 +156,10 @@ def _set_device(self, device):
def _serialize_deserialize(self, trt_engine):
self.runtime = trt.Runtime(TRT_LOGGER)
serialized_engine = trt_engine.serialize()

with open(self.engine_path,'wb') as w:
w.write(serialized_engine)

del self.parser # Parser no longer needed for ownership of plugins
trt_engine = self.runtime.deserialize_cuda_engine(
serialized_engine)
Expand Down