wehub-resource-sync
86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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.
|
|
#
|
|
import os
|
|
import onnx
|
|
import pytest
|
|
from polygraphy import util, mod
|
|
from tests.helper import TEST_DIR, poly_run, is_file_non_empty
|
|
from tests.models.meta import ONNX_MODELS_PATH
|
|
|
|
class TestModelEngine:
|
|
def test_load_onnx_model(self):
|
|
# Load an onnx model and run inference
|
|
assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec"])
|
|
|
|
def test_save_load_engine(self):
|
|
# Save an engine and check if file exists. Run inference with saved engine
|
|
with util.NamedTemporaryFile() as outpath:
|
|
assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--save-engine", outpath.name])
|
|
assert is_file_non_empty(outpath.name)
|
|
assert poly_run(["--trtexec", outpath.name, "--model-type=engine"])
|
|
|
|
class TestProfileShapes:
|
|
def test_input_shape(self):
|
|
assert poly_run([ONNX_MODELS_PATH['dynamic_identity'],
|
|
"--trtexec",
|
|
"--input-shapes", "X:[1,2,3,3]",
|
|
])
|
|
|
|
def test_explicit_profile(self):
|
|
assert poly_run([ ONNX_MODELS_PATH['dynamic_identity'],
|
|
"--trtexec",
|
|
"--input-shapes", "X:[1,2,1,1]",
|
|
"--trt-min-shapes", "X:[1,2,1,1]",
|
|
"--trt-opt-shapes", "X:[1,2,1,1]",
|
|
"--trt-max-shapes", "X:[1,2,1,1]",
|
|
])
|
|
|
|
def test_explicit_profile_opt_runtime_shapes_differ(self):
|
|
assert poly_run([ONNX_MODELS_PATH['dynamic_identity'],
|
|
"--trtexec",
|
|
"--input-shapes", "X:[1,2,2,2]",
|
|
"--trt-min-shapes", "X:[1,2,1,1]",
|
|
"--trt-opt-shapes", "X:[1,2,3,3]",
|
|
"--trt-max-shapes", "X:[1,2,4,4]",
|
|
])
|
|
|
|
class TestParameterMapping:
|
|
def test_args(self):
|
|
assert poly_run([ONNX_MODELS_PATH['identity'],
|
|
"--trtexec",
|
|
"--fp16", "--int8", "--best", "--tf32",
|
|
"--use-cuda-graph", "--no-data-transfers",
|
|
])
|
|
|
|
def test_kwargs(self):
|
|
assert poly_run([ONNX_MODELS_PATH['identity'],
|
|
"--trtexec",
|
|
"--memory-pool-limit=workspace:1M", "--streams=2", "--device=0",
|
|
"--avg-runs=1", "--min-timing=1", "--avg-timing=3",
|
|
"--trtexec-iterations=1", "--trtexec-warmup=1", "--duration=1",
|
|
])
|
|
|
|
|
|
class TestAccuracyComparision:
|
|
def test_trtexec_trt(self):
|
|
# Compare inference results between trtexec and trt
|
|
assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--trt"])
|
|
|
|
def test_trtexec_onnxrt(self):
|
|
# Compare inference results between trtexec and onnxrt
|
|
assert poly_run([ONNX_MODELS_PATH['identity'], "--trtexec", "--onnxrt"])
|