File size: 2,016 Bytes
b386992 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
#
# 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.
try:
import tensorrt as trt
from polygraphy.backend.trt import CreateConfig, Profile, engine_from_network, network_from_onnx_path, save_engine
HAVE_TRT = True
except (ImportError, ModuleNotFoundError):
HAVE_TRT = False
def build_engine(
onnx_path,
output_path,
fp16,
input_profile=None,
enable_refit=False,
enable_preview=False,
timing_cache=None,
workspace_size=0,
int8=False,
builder_optimization_level=None,
):
print(f"Building TensorRT engine for {onnx_path}: {output_path}")
p = Profile()
if input_profile:
for name, dims in input_profile.items():
assert len(dims) == 3
p.add(name, min=dims[0], opt=dims[1], max=dims[2])
preview_features = None
config_kwargs = {}
if workspace_size > 0:
config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size}
engine = engine_from_network(
network_from_onnx_path(onnx_path),
config=CreateConfig(
fp16=fp16,
refittable=enable_refit,
profiles=[p],
preview_features=preview_features,
load_timing_cache=timing_cache,
int8=int8,
builder_optimization_level=builder_optimization_level,
**config_kwargs,
),
save_timing_cache=timing_cache,
)
save_engine(engine, path=output_path)
|