wehub-resource-sync
409 lines
18 KiB
Python
409 lines
18 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 threading
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from polygraphy import config, cuda, mod
|
|
from polygraphy.backend.trt import (
|
|
EngineFromNetwork,
|
|
NetworkFromOnnxBytes,
|
|
Profile,
|
|
TrtRunner,
|
|
engine_from_network,
|
|
network_from_onnx_bytes,
|
|
)
|
|
from polygraphy.backend.trt.runner import _get_array_on_cpu
|
|
from polygraphy.exception import PolygraphyException
|
|
from polygraphy.logger import G_LOGGER
|
|
from tests.models.meta import ONNX_MODELS
|
|
|
|
# Import CreateConfigRTX conditionally for TensorRT-RTX builds
|
|
if config.USE_TENSORRT_RTX:
|
|
import tensorrt_rtx as trt
|
|
from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig
|
|
else:
|
|
import tensorrt as trt
|
|
from polygraphy.backend.trt import CreateConfig
|
|
|
|
|
|
class TestLoggerCallbacks:
|
|
@pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
|
|
def test_set_severity(self, sev):
|
|
G_LOGGER.module_severity = sev
|
|
|
|
|
|
@pytest.fixture(scope="class")
|
|
def nonzero_engine():
|
|
model = ONNX_MODELS["nonzero"]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
return engine_from_network(network_loader)
|
|
|
|
|
|
@pytest.fixture()
|
|
def identity_engine():
|
|
model = ONNX_MODELS["identity"]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
return engine_from_network(network_loader)
|
|
|
|
|
|
@pytest.fixture()
|
|
def reducable_engine():
|
|
model = ONNX_MODELS["reducable"]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
return engine_from_network(network_loader)
|
|
|
|
|
|
class TestTrtRunner:
|
|
def test_can_name_runner(self):
|
|
NAME = "runner"
|
|
runner = TrtRunner(None, name=NAME)
|
|
assert runner.name == NAME
|
|
|
|
def test_basic(self, identity_engine):
|
|
with TrtRunner(identity_engine) as runner:
|
|
assert runner.optimization_profile is None
|
|
assert runner.is_active
|
|
ONNX_MODELS["identity"].check_runner(runner)
|
|
assert runner.last_inference_time() is not None
|
|
assert not runner.is_active
|
|
|
|
@pytest.mark.serial
|
|
@pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX has different warning output behavior")
|
|
def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods, identity_engine):
|
|
runner = TrtRunner(identity_engine)
|
|
check_warnings_on_runner_impl_methods(runner)
|
|
|
|
@pytest.mark.parametrize(
|
|
"inp, expected",
|
|
[
|
|
([1, 0, 1, 1], [[0, 2, 3]]),
|
|
([1, 0, 0, 1], [[0, 3]]),
|
|
([0, 0, 0, 1], [[3]]),
|
|
],
|
|
)
|
|
@pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support data dependent shapes")
|
|
def test_data_dependent_shapes(self, nonzero_engine, inp, expected):
|
|
with TrtRunner(nonzero_engine) as runner:
|
|
outputs = runner.infer(
|
|
{
|
|
"input": np.array(
|
|
inp,
|
|
dtype=(np.int32 if mod.version(trt.__version__) < mod.version("9.0") else np.int64),
|
|
)
|
|
}
|
|
)
|
|
assert np.array_equal(outputs["nonzero_out_0"], np.array(expected, dtype=np.int32))
|
|
|
|
@pytest.mark.parametrize("copy_outputs_to_host", [True, False])
|
|
@pytest.mark.parametrize("device", ["cpu", "cuda"])
|
|
def test_torch_tensors(self, copy_outputs_to_host, identity_engine, device):
|
|
with TrtRunner(identity_engine) as runner:
|
|
arr = torch.ones([1, 1, 2, 2], dtype=torch.float32, device=device)
|
|
outputs = runner.infer({"x": arr}, copy_outputs_to_host=copy_outputs_to_host)
|
|
assert all(isinstance(t, torch.Tensor) for t in outputs.values())
|
|
|
|
assert torch.equal(outputs["y"].to("cpu"), arr.to("cpu"))
|
|
|
|
assert outputs["y"].device.type == ("cpu" if copy_outputs_to_host else "cuda")
|
|
|
|
def test_context(self, identity_engine):
|
|
with TrtRunner(identity_engine.create_execution_context) as runner:
|
|
ONNX_MODELS["identity"].check_runner(runner)
|
|
|
|
def test_device_buffer_order_matches_bindings(self, reducable_engine):
|
|
with TrtRunner(reducable_engine) as runner:
|
|
dev_buf_order = list(runner.device_input_buffers.keys())
|
|
for binding, dev_buf_name in zip(reducable_engine, dev_buf_order):
|
|
assert binding == dev_buf_name
|
|
|
|
def test_shape_output(self):
|
|
model = ONNX_MODELS["reshape"]
|
|
engine = engine_from_network(NetworkFromOnnxBytes(model.loader))
|
|
with engine, TrtRunner(engine.create_execution_context) as runner:
|
|
model.check_runner(runner)
|
|
|
|
def test_multithreaded_runners_from_engine(self, identity_engine):
|
|
with TrtRunner(identity_engine) as runner0, TrtRunner(identity_engine) as runner1:
|
|
t1 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner0,))
|
|
t2 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner1,))
|
|
t1.start()
|
|
t2.start()
|
|
t1.join()
|
|
t2.join()
|
|
|
|
@pytest.mark.parametrize("use_optimization_profile", [True, False])
|
|
@pytest.mark.skipif(
|
|
not config.USE_TENSORRT_RTX
|
|
and mod.version(trt.__version__) >= mod.version("8.6")
|
|
and mod.version(trt.__version__) < mod.version("8.7"),
|
|
reason="Bug in TRT 8.6",
|
|
)
|
|
def test_multiple_profiles(self, use_optimization_profile):
|
|
model = ONNX_MODELS["dynamic_identity"]
|
|
profile0_shapes = [
|
|
(1, 2, 1, 1),
|
|
(1, 2, 1, 1),
|
|
(1, 2, 1, 1),
|
|
] # Use min==opt==max to fix shapes in the engine.
|
|
profile1_shapes = [(1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)]
|
|
profile2_shapes = [(1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
profiles = [
|
|
Profile().add("X", *profile0_shapes),
|
|
Profile().add("X", *profile1_shapes),
|
|
Profile().add("X", *profile2_shapes),
|
|
]
|
|
config_loader = CreateConfig(profiles=profiles)
|
|
engine = engine_from_network(network_loader, config_loader)
|
|
context = engine.create_execution_context()
|
|
|
|
for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
|
|
with TrtRunner(
|
|
context,
|
|
optimization_profile=index if use_optimization_profile else None,
|
|
) as runner:
|
|
if not use_optimization_profile:
|
|
runner.set_profile(index)
|
|
|
|
assert runner.context.active_optimization_profile == index
|
|
for shape in shapes:
|
|
model.check_runner(runner, {"X": shape})
|
|
|
|
@pytest.mark.skipif(
|
|
not config.USE_TENSORRT_RTX and mod.version(trt.__version__) < mod.version("10.0"),
|
|
reason="Feature not present before 10.0",
|
|
)
|
|
@pytest.mark.parametrize("allocation_strategy", [None, "static", "profile", "runtime"])
|
|
def test_allocation_strategies(self, allocation_strategy):
|
|
if config.USE_TENSORRT_RTX and allocation_strategy == "runtime":
|
|
pytest.skip("TensorRT-RTX issues with runtime allocation strategy")
|
|
|
|
model = ONNX_MODELS["residual_block"]
|
|
profile0_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (1, 3, 224, 224)]
|
|
profile1_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224)]
|
|
profile2_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (4, 3, 224, 224)]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
profiles = [
|
|
Profile().add("gpu_0/data_0", *profile0_shapes),
|
|
Profile().add("gpu_0/data_0", *profile1_shapes),
|
|
Profile().add("gpu_0/data_0", *profile2_shapes),
|
|
]
|
|
config_loader = CreateConfig(profiles=profiles)
|
|
engine = engine_from_network(network_loader, config_loader)
|
|
|
|
for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]):
|
|
with TrtRunner(
|
|
engine,
|
|
optimization_profile=index,
|
|
allocation_strategy=allocation_strategy,
|
|
) as runner:
|
|
for shape in shapes:
|
|
model.check_runner(runner, {"gpu_0/data_0": shape})
|
|
|
|
def test_empty_tensor_with_dynamic_input_shape_tensor(self):
|
|
model = ONNX_MODELS["empty_tensor_expand"]
|
|
shapes = [(1, 2, 0, 3, 0), (2, 2, 0, 3, 0), (4, 2, 0, 3, 0)]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
profiles = [Profile().add("new_shape", *shapes)]
|
|
config_loader = CreateConfig(profiles=profiles)
|
|
|
|
with TrtRunner(EngineFromNetwork(network_loader, config_loader)) as runner:
|
|
for shape in shapes:
|
|
model.check_runner(runner, {"new_shape": shape})
|
|
|
|
@pytest.mark.parametrize(
|
|
"names, err",
|
|
[
|
|
(["fake-input", "x"], "Extra inputs in"),
|
|
(["fake-input"], "The following inputs were not found"),
|
|
([], "The following inputs were not found"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("module", [torch, np])
|
|
def test_error_on_wrong_name_feed_dict(self, names, err, identity_engine, module):
|
|
with TrtRunner(identity_engine) as runner:
|
|
with pytest.raises(PolygraphyException, match=err):
|
|
runner.infer({name: module.ones((1, 1, 2, 2), dtype=module.float32) for name in names})
|
|
|
|
@pytest.mark.parametrize("module", [torch, np])
|
|
def test_error_on_wrong_dtype_feed_dict(self, identity_engine, module):
|
|
with TrtRunner(identity_engine) as runner:
|
|
with pytest.raises(PolygraphyException, match="unexpected dtype."):
|
|
runner.infer({"x": module.ones((1, 1, 2, 2), dtype=module.int32)})
|
|
|
|
@pytest.mark.parametrize("module", [torch, np])
|
|
def test_error_on_wrong_shape_feed_dict(self, identity_engine, module):
|
|
with TrtRunner(identity_engine) as runner:
|
|
with pytest.raises(PolygraphyException, match="incompatible shape."):
|
|
runner.infer({"x": module.ones((1, 1, 3, 2), dtype=module.float32)})
|
|
|
|
@pytest.mark.parametrize("use_view", [True, False]) # We should be able to use DeviceArray in place of DeviceView
|
|
def test_device_views(self, use_view, reducable_engine):
|
|
with TrtRunner(reducable_engine) as runner, cuda.DeviceArray((1,), dtype=np.float32) as x:
|
|
x.copy_from(np.ones((1,), dtype=np.float32))
|
|
outputs = runner.infer(
|
|
{
|
|
"X0": x.view() if use_view else x,
|
|
"Y0": np.ones((1,), dtype=np.float32),
|
|
}
|
|
)
|
|
assert outputs["identity_out_6"][0] == 2
|
|
assert outputs["identity_out_8"][0] == 2
|
|
|
|
def test_no_output_copy(self, identity_engine):
|
|
with TrtRunner(identity_engine) as runner:
|
|
inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
|
|
outputs = runner.infer({"x": inp}, copy_outputs_to_host=False)
|
|
assert isinstance(outputs["y"], cuda.DeviceView)
|
|
assert np.array_equal(outputs["y"].numpy(), inp)
|
|
|
|
def test_subsequent_infers_with_different_input_types(self, identity_engine):
|
|
with TrtRunner(identity_engine) as runner:
|
|
inp = np.ones(shape=(1, 1, 2, 2), dtype=np.float32)
|
|
|
|
def check(outputs):
|
|
assert np.all(outputs["y"] == inp)
|
|
|
|
check(runner.infer({"x": inp}))
|
|
check(runner.infer({"x": cuda.DeviceArray(shape=inp.shape, dtype=inp.dtype).copy_from(inp)}))
|
|
|
|
torch_outputs = runner.infer({"x": torch.from_numpy(inp)})
|
|
check({name: out.numpy() for name, out in torch_outputs.items()})
|
|
check(runner.infer({"x": inp}))
|
|
|
|
@pytest.mark.parametrize("use_view", [True, False]) # We should be able to use DeviceArray in place of DeviceView
|
|
def test_device_view_dynamic_shapes(self, use_view):
|
|
model = ONNX_MODELS["dynamic_identity"]
|
|
profiles = [
|
|
Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
|
|
]
|
|
runner = TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader), CreateConfig(profiles=profiles)))
|
|
with runner, cuda.DeviceArray(shape=(1, 2, 3, 3), dtype=np.float32) as arr:
|
|
inp = np.random.random_sample(size=(1, 2, 3, 3)).astype(np.float32)
|
|
arr.copy_from(inp)
|
|
outputs = runner.infer({"X": (cuda.DeviceView(arr.ptr, arr.shape, arr.dtype) if use_view else arr)})
|
|
assert np.all(outputs["Y"] == inp)
|
|
assert outputs["Y"].shape == (1, 2, 3, 3)
|
|
|
|
def test_cannot_use_device_view_shape_tensor(self):
|
|
model = ONNX_MODELS["empty_tensor_expand"]
|
|
with TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader))) as runner, cuda.DeviceArray(
|
|
shape=(5,),
|
|
dtype=(
|
|
np.int32
|
|
if mod.version(trt.__version__) < mod.version("9.0") and not config.USE_TENSORRT_RTX
|
|
else np.int64
|
|
),
|
|
) as arr:
|
|
with pytest.raises(PolygraphyException, match="it must reside in host memory"):
|
|
runner.infer({"data": np.ones((2, 0, 3, 0), dtype=np.float32), "new_shape": arr})
|
|
|
|
@pytest.mark.parametrize("hwc_input", [True, False], ids=["hwc_input", "chw_input"])
|
|
@pytest.mark.parametrize("hwc_output", [True, False], ids=["hwc_output", "chw_output"])
|
|
@pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support custom I/O format networks")
|
|
def test_infer_chw_format(self, hwc_input, hwc_output):
|
|
model = ONNX_MODELS["identity_multi_ch"]
|
|
inp_shape = model.input_metadata["x"].shape
|
|
builder, network, parser = network_from_onnx_bytes(model.loader)
|
|
|
|
formats = 1 << int(trt.TensorFormat.HWC)
|
|
if hwc_input:
|
|
network.get_input(0).allowed_formats = formats
|
|
if hwc_output:
|
|
network.get_output(0).allowed_formats = formats
|
|
|
|
engine = engine_from_network((builder, network))
|
|
|
|
with TrtRunner(engine) as runner:
|
|
inp = np.random.normal(size=(inp_shape)).astype(np.float32)
|
|
if hwc_input:
|
|
inp = inp.transpose(0, 2, 3, 1)
|
|
|
|
outputs = runner.infer({"x": inp})
|
|
if hwc_input == hwc_output: # output in CHW/HWC format and similarly shaped
|
|
assert np.allclose(outputs["y"], inp)
|
|
elif not hwc_input and hwc_output: # output in HWC format and shaped (N, H, W, C)
|
|
assert np.allclose(outputs["y"].transpose(0, 3, 1, 2), inp)
|
|
else: # hwc_input and not hwc_output: output in CHW format and shaped (N, C, H, W)
|
|
assert np.allclose(outputs["y"].transpose(0, 2, 3, 1), inp)
|
|
|
|
@pytest.mark.parametrize("use_torch", [True, False])
|
|
def test_get_array_on_cpu(self, use_torch):
|
|
shape = (4,)
|
|
with cuda.DeviceArray.raw(shape) as arr:
|
|
host_buffers = {}
|
|
stream = cuda.Stream()
|
|
host_arr = _get_array_on_cpu(arr, "test", host_buffers, stream, arr.nbytes, use_torch)
|
|
|
|
if use_torch:
|
|
assert isinstance(host_arr, torch.Tensor)
|
|
else:
|
|
assert isinstance(host_arr, np.ndarray)
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX,
|
|
reason="Feature not present before 10.0",
|
|
)
|
|
@pytest.mark.parametrize("budget", [None, -2, -1, 0, 0.5, 0.99, 1.0, 1000, np.inf])
|
|
def test_weight_streaming(self, budget):
|
|
model = ONNX_MODELS["matmul_2layer"]
|
|
network_loader = NetworkFromOnnxBytes(model.loader, strongly_typed=True)
|
|
config_loader = CreateConfig(weight_streaming=True)
|
|
engine = engine_from_network(network_loader, config_loader)
|
|
|
|
if budget == np.inf:
|
|
# set to max size - 1
|
|
budget = engine.streamable_weights_size - 1
|
|
|
|
kwargs = {"weight_streaming_budget": None, "weight_streaming_percent": None}
|
|
if budget is not None:
|
|
if 0 < budget <= 1:
|
|
kwargs["weight_streaming_percent"] = budget * 100
|
|
else:
|
|
kwargs["weight_streaming_budget"] = int(budget)
|
|
|
|
with TrtRunner(engine, optimization_profile=0, **kwargs) as runner:
|
|
model.check_runner(runner)
|
|
|
|
@pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
|
|
def test_compute_capabilities_engine_building(self):
|
|
"""Test compute capabilities integration with engine building"""
|
|
model = ONNX_MODELS["identity"]
|
|
network_loader = NetworkFromOnnxBytes(model.loader)
|
|
|
|
# Test --use-gpu flag
|
|
config_loader = CreateConfig(use_gpu=True)
|
|
engine = engine_from_network(network_loader, config_loader)
|
|
with TrtRunner(engine) as runner:
|
|
model.check_runner(runner)
|
|
|
|
# Test --compute-capabilities flag
|
|
config_loader = CreateConfig(compute_capabilities=[(7, 5), (8, 0), (8, 6)])
|
|
engine = engine_from_network(network_loader, config_loader)
|
|
with TrtRunner(engine) as runner:
|
|
model.check_runner(runner)
|
|
|
|
@pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled")
|
|
def test_compute_capabilities_mutual_exclusion(self):
|
|
"""Test that use_gpu and compute_capabilities are mutually exclusive"""
|
|
# Test mutual exclusion - should raise an exception
|
|
with pytest.raises(PolygraphyException, match="use_gpu and compute_capabilities are mutually exclusive"):
|
|
CreateConfig(use_gpu=True, compute_capabilities=[(7, 5)])
|