wehub-resource-sync
543 lines
22 KiB
Python
543 lines
22 KiB
Python
import contextlib
|
|
import os
|
|
import tempfile
|
|
|
|
import pytest
|
|
|
|
from polygraphy import mod, util
|
|
from polygraphy.backend.trt import (
|
|
Calibrator,
|
|
Profile,
|
|
network_from_onnx_bytes,
|
|
postprocess_config,
|
|
)
|
|
from polygraphy.common.struct import BoundedShape
|
|
from polygraphy.comparator import DataLoader
|
|
from polygraphy.datatype import DataType
|
|
from polygraphy import config as polygraphy_config
|
|
from tests.helper import has_dla
|
|
from tests.models.meta import ONNX_MODELS
|
|
|
|
# Import CreateConfigRTX conditionally for TensorRT-RTX builds
|
|
if polygraphy_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
|
|
|
|
@pytest.fixture(scope="session")
|
|
def identity_builder_network():
|
|
builder, network, parser = network_from_onnx_bytes(ONNX_MODELS["identity"].loader)
|
|
with builder, network, parser:
|
|
yield builder, network
|
|
|
|
|
|
class TestCreateConfig:
|
|
def test_defaults(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig()
|
|
assert loader.timing_cache_path is None
|
|
|
|
with loader(builder, network) as config:
|
|
assert not config.get_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE)
|
|
with contextlib.suppress(AttributeError):
|
|
if polygraphy_config.USE_TENSORRT_RTX:
|
|
assert config.get_flag(trt.BuilderFlag.TF32)
|
|
else:
|
|
assert not config.get_flag(trt.BuilderFlag.TF32)
|
|
with contextlib.suppress(AttributeError):
|
|
assert not config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
|
|
assert not config.get_flag(trt.BuilderFlag.FP16)
|
|
assert not config.get_flag(trt.BuilderFlag.INT8)
|
|
if mod.version(trt.__version__) >= mod.version("8.7"):
|
|
assert not config.get_flag(trt.BuilderFlag.BF16)
|
|
if mod.version(trt.__version__) >= mod.version("8.6"):
|
|
assert not config.get_flag(trt.BuilderFlag.FP8)
|
|
assert not config.get_flag(trt.BuilderFlag.VERSION_COMPATIBLE)
|
|
assert not config.get_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME)
|
|
if not polygraphy_config.USE_TENSORRT_RTX:
|
|
assert (
|
|
config.hardware_compatibility_level
|
|
== trt.HardwareCompatibilityLevel.NONE
|
|
)
|
|
if mod.version(trt.__version__) >= mod.version("10.2") and not polygraphy_config.USE_TENSORRT_RTX:
|
|
assert (
|
|
config.runtime_platform
|
|
== trt.RuntimePlatform.SAME_AS_BUILD
|
|
)
|
|
assert config.num_optimization_profiles == 1
|
|
if not polygraphy_config.USE_TENSORRT_RTX:
|
|
assert config.int8_calibrator is None
|
|
with contextlib.suppress(AttributeError):
|
|
if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
|
|
assert config.get_tactic_sources() == 24
|
|
elif mod.version(trt.__version__) >= mod.version("8.7"):
|
|
assert config.get_tactic_sources() == 29
|
|
elif mod.version(trt.__version__) >= mod.version("8.5"):
|
|
assert config.get_tactic_sources() == 31
|
|
if mod.version(trt.__version__) >= mod.version("8.7"):
|
|
assert not config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS)
|
|
if mod.version(trt.__version__) >= mod.version("8.7"):
|
|
assert not config.get_flag(trt.BuilderFlag.DISABLE_COMPILATION_CACHE)
|
|
with contextlib.suppress(AttributeError):
|
|
assert not config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
|
|
with contextlib.suppress(AttributeError):
|
|
if polygraphy_config.USE_TENSORRT_RTX:
|
|
assert config.engine_capability == trt.EngineCapability.STANDARD
|
|
else:
|
|
assert config.engine_capability == trt.EngineCapability.DEFAULT
|
|
with contextlib.suppress(AttributeError):
|
|
assert not config.get_flag(trt.BuilderFlag.DIRECT_IO)
|
|
|
|
@pytest.mark.parametrize(
|
|
"engine_capability",
|
|
[
|
|
trt.EngineCapability.STANDARD,
|
|
trt.EngineCapability.SAFETY,
|
|
trt.EngineCapability.DLA_STANDALONE,
|
|
],
|
|
)
|
|
def test_engine_capability(self, identity_builder_network, engine_capability):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(engine_capability=engine_capability)
|
|
with loader(builder, network) as config:
|
|
assert config.engine_capability == engine_capability
|
|
|
|
@pytest.mark.parametrize("flag", ["obey", "prefer", None])
|
|
def test_precision_constraints(self, identity_builder_network, flag):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(precision_constraints=flag)
|
|
with loader(builder, network) as config:
|
|
obey_set = config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
|
|
prefer_set = config.get_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
|
|
if flag == "obey":
|
|
assert obey_set and not prefer_set
|
|
elif flag == "prefer":
|
|
assert not obey_set and prefer_set
|
|
else:
|
|
assert not obey_set and not prefer_set
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="Unsupported before TRT 8.6",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"kwargs, expected_flag",
|
|
[
|
|
({"version_compatible": True}, "VERSION_COMPATIBLE"),
|
|
(
|
|
{"version_compatible": True, "exclude_lean_runtime": True},
|
|
"EXCLUDE_LEAN_RUNTIME",
|
|
),
|
|
],
|
|
)
|
|
def test_version_compatibility_flags(
|
|
self, identity_builder_network, kwargs, expected_flag
|
|
):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(**kwargs)
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(getattr(trt.BuilderFlag, expected_flag))
|
|
|
|
def test_direct_io(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(direct_io=True)
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(trt.BuilderFlag.DIRECT_IO)
|
|
|
|
@pytest.mark.parametrize("flag", [True, False])
|
|
def test_restricted(self, identity_builder_network, flag):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(restricted=flag)
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(trt.BuilderFlag.SAFETY_SCOPE) == flag
|
|
|
|
@pytest.mark.parametrize(
|
|
"arg_name, flag_type",
|
|
[
|
|
("refittable", trt.BuilderFlag.REFIT),
|
|
]
|
|
+ (
|
|
[
|
|
(
|
|
"disable_compilation_cache",
|
|
trt.BuilderFlag.DISABLE_COMPILATION_CACHE,
|
|
),
|
|
]
|
|
if mod.version(trt.__version__) >= mod.version("9.0")
|
|
else []
|
|
)
|
|
+ (
|
|
[
|
|
("strip_plan", trt.BuilderFlag.STRIP_PLAN),
|
|
]
|
|
if mod.version(trt.__version__) >= mod.version("10.0")
|
|
else []
|
|
)
|
|
+ (
|
|
[
|
|
("fp16", trt.BuilderFlag.FP16),
|
|
("int8", trt.BuilderFlag.INT8),
|
|
("allow_gpu_fallback", trt.BuilderFlag.GPU_FALLBACK),
|
|
("tf32", trt.BuilderFlag.TF32),
|
|
]
|
|
+ (
|
|
[
|
|
("bf16", trt.BuilderFlag.BF16),
|
|
]
|
|
if mod.version(trt.__version__) >= mod.version("8.7")
|
|
else []
|
|
)
|
|
+ (
|
|
[
|
|
("fp8", trt.BuilderFlag.FP8),
|
|
]
|
|
if mod.version(trt.__version__) >= mod.version("8.6")
|
|
else []
|
|
)
|
|
if not polygraphy_config.USE_TENSORRT_RTX
|
|
else []
|
|
),
|
|
)
|
|
@pytest.mark.parametrize("value", [True, False])
|
|
def test_flags(self, identity_builder_network, arg_name, flag_type, value):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(**{arg_name: value})
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(flag_type) == value
|
|
|
|
@pytest.mark.parametrize("flag", [True, False])
|
|
def test_sparse_weights(self, identity_builder_network, flag):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(sparse_weights=flag)
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS) == flag
|
|
|
|
@pytest.mark.skipif(
|
|
polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="TensorRT-RTX does not support DLA"
|
|
)
|
|
def test_use_dla(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(use_dla=True)
|
|
with loader(builder, network) as config:
|
|
assert config.default_device_type == trt.DeviceType.DLA
|
|
if has_dla():
|
|
assert config.DLA_core == 0
|
|
|
|
with contextlib.suppress(AttributeError):
|
|
TACTIC_SOURCES_CASES = [
|
|
(None, 31), # By default, all sources are enabled.
|
|
([], 0),
|
|
([trt.TacticSource.CUBLAS], 1),
|
|
([trt.TacticSource.CUBLAS_LT], 2),
|
|
([trt.TacticSource.CUDNN], 4),
|
|
([trt.TacticSource.CUBLAS, trt.TacticSource.CUBLAS_LT], 3),
|
|
([trt.TacticSource.CUBLAS, trt.TacticSource.CUDNN], 5),
|
|
([trt.TacticSource.CUBLAS_LT, trt.TacticSource.CUDNN], 6),
|
|
(
|
|
[
|
|
trt.TacticSource.CUDNN,
|
|
trt.TacticSource.CUBLAS,
|
|
trt.TacticSource.CUBLAS_LT,
|
|
],
|
|
7,
|
|
),
|
|
(
|
|
[
|
|
trt.TacticSource.CUDNN,
|
|
trt.TacticSource.CUBLAS,
|
|
trt.TacticSource.CUBLAS_LT,
|
|
trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
|
|
],
|
|
15,
|
|
),
|
|
(
|
|
[
|
|
trt.TacticSource.CUDNN,
|
|
trt.TacticSource.CUBLAS,
|
|
trt.TacticSource.CUBLAS_LT,
|
|
trt.TacticSource.EDGE_MASK_CONVOLUTIONS,
|
|
trt.TacticSource.JIT_CONVOLUTIONS,
|
|
],
|
|
31,
|
|
),
|
|
]
|
|
|
|
if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX:
|
|
TACTIC_SOURCES_CASES[0] = (None, 24)
|
|
elif mod.version(trt.__version__) >= mod.version("8.7"):
|
|
TACTIC_SOURCES_CASES[0] = (None, 29)
|
|
|
|
@pytest.mark.parametrize("sources, expected", TACTIC_SOURCES_CASES)
|
|
def test_tactic_sources(self, identity_builder_network, sources, expected):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(tactic_sources=sources)
|
|
with loader(builder, network) as config:
|
|
assert config.get_tactic_sources() == expected
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("8.7") and not polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="API was added in TRT 8.7",
|
|
)
|
|
@pytest.mark.parametrize("flag", [True, False])
|
|
def test_error_on_timing_cache_miss(self, identity_builder_network, flag):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(error_on_timing_cache_miss=flag)
|
|
with loader(builder, network) as config:
|
|
assert config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS) == flag
|
|
|
|
@pytest.mark.skipif(
|
|
polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="TensorRT-RTX does not support calibrators"
|
|
)
|
|
def test_calibrator_metadata_set(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
calibrator = Calibrator(DataLoader())
|
|
loader = CreateConfig(int8=True, calibrator=calibrator)
|
|
with loader(builder, network) as config:
|
|
assert config.int8_calibrator
|
|
assert "x" in calibrator.data_loader.input_metadata
|
|
meta = calibrator.data_loader.input_metadata["x"]
|
|
assert meta.shape == BoundedShape((1, 1, 2, 2))
|
|
assert meta.dtype == DataType.FLOAT32
|
|
|
|
def test_multiple_profiles(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
profiles = [
|
|
Profile().add("x", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)),
|
|
Profile().add("x", (1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)),
|
|
]
|
|
loader = CreateConfig(profiles=profiles)
|
|
with loader(builder, network) as config:
|
|
assert config.num_optimization_profiles == 2
|
|
|
|
@pytest.mark.parametrize("path_mode", [True, False], ids=["path", "file-like"])
|
|
def test_timing_cache(self, identity_builder_network, path_mode):
|
|
builder, network = identity_builder_network
|
|
with util.NamedTemporaryFile() as cache:
|
|
loader = CreateConfig(load_timing_cache=cache.name if path_mode else cache)
|
|
with loader(builder, network) as config:
|
|
assert config.get_timing_cache()
|
|
|
|
def test_fall_back_to_empty_timing_cache(self, identity_builder_network):
|
|
"""Tests that passing in a nonexistent timing cache path is non-fatal"""
|
|
builder, network = identity_builder_network
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
cache_name = os.path.join(tmpdir, "casper")
|
|
loader = CreateConfig(load_timing_cache=cache_name)
|
|
with loader(builder, network) as config:
|
|
assert config.get_timing_cache()
|
|
|
|
def test_empty_timing_cache_when_default(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig()
|
|
with loader(builder, network) as config:
|
|
cache = config.get_timing_cache()
|
|
with cache.serialize() as buffer:
|
|
cache_size = len(bytes(buffer))
|
|
|
|
cache.reset()
|
|
with cache.serialize() as buffer:
|
|
new_cache_size = len(bytes(buffer))
|
|
assert cache_size == new_cache_size
|
|
|
|
def test_profiling_verbosity(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
expected = trt.ProfilingVerbosity.NONE
|
|
loader = CreateConfig(profiling_verbosity=expected)
|
|
with loader(builder, network) as config:
|
|
assert config.profiling_verbosity == expected
|
|
|
|
with contextlib.suppress(AttributeError):
|
|
POOL_LIMITS = [
|
|
{trt.MemoryPoolType.WORKSPACE: 25},
|
|
{trt.MemoryPoolType.DLA_MANAGED_SRAM: 25},
|
|
{trt.MemoryPoolType.DLA_LOCAL_DRAM: 25},
|
|
{trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25},
|
|
# Multiple limits
|
|
{
|
|
trt.MemoryPoolType.DLA_LOCAL_DRAM: 20,
|
|
trt.MemoryPoolType.DLA_GLOBAL_DRAM: 25,
|
|
trt.MemoryPoolType.WORKSPACE: 39,
|
|
},
|
|
]
|
|
|
|
# @pytest.mark.skipif(
|
|
# config.USE_TENSORRT_RTX,
|
|
# reason="TensorRT-RTX does not support DLA memory pools"
|
|
# )
|
|
@pytest.mark.parametrize("pool_limits", POOL_LIMITS)
|
|
def test_memory_pool_limits(self, pool_limits, identity_builder_network):
|
|
if any("dla" in key.name.lower() for key in pool_limits) and not has_dla():
|
|
pytest.skip("DLA is not available on this system")
|
|
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(memory_pool_limits=pool_limits)
|
|
with loader(builder, network) as config:
|
|
for pool_type, pool_size in pool_limits.items():
|
|
assert config.get_memory_pool_limit(pool_type) == pool_size
|
|
|
|
@pytest.mark.parametrize(
|
|
"preview_features",
|
|
[
|
|
[trt.PreviewFeature.PROFILE_SHARING_0806]
|
|
if mod.version(trt.__version__) >= mod.version("10.0")
|
|
else (
|
|
[trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
|
|
if polygraphy_config.USE_TENSORRT_RTX
|
|
else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
|
|
),
|
|
],
|
|
)
|
|
def test_preview_features(self, identity_builder_network, preview_features):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(preview_features=preview_features)
|
|
with loader(builder, network) as config:
|
|
# Check that only the enabled preview features are on.
|
|
for pf in trt.PreviewFeature.__members__.values():
|
|
expected = pf in preview_features
|
|
# TensorRT-RTX enables PROFILE_SHARING_0806 by default and can't be disabled
|
|
if polygraphy_config.USE_TENSORRT_RTX and pf == trt.PreviewFeature.PROFILE_SHARING_0806:
|
|
expected = True
|
|
assert config.get_preview_feature(pf) == expected
|
|
|
|
@pytest.mark.skipif(
|
|
polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="TensorRT-RTX does not support quantization_flag API"
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"quantization_flags",
|
|
[
|
|
[],
|
|
[trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION],
|
|
],
|
|
)
|
|
def test_quantization_flags(self, identity_builder_network, quantization_flags):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(quantization_flags=quantization_flags)
|
|
with loader(builder, network) as config:
|
|
# Check that only the enabled quantization flags are on.
|
|
for qf in trt.QuantizationFlag.__members__.values():
|
|
assert config.get_quantization_flag(qf) == (qf in quantization_flags)
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="Unsupported for TRT versions prior to 8.6",
|
|
)
|
|
@pytest.mark.parametrize("level", range(6))
|
|
def test_builder_optimization_level(self, identity_builder_network, level):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(builder_optimization_level=level)
|
|
with loader(builder, network) as config:
|
|
assert config.builder_optimization_level == level
|
|
|
|
if mod.version(trt.__version__) >= mod.version("8.6"):
|
|
|
|
@pytest.mark.parametrize(
|
|
"level",
|
|
[
|
|
trt.HardwareCompatibilityLevel.NONE,
|
|
trt.HardwareCompatibilityLevel.AMPERE_PLUS,
|
|
],
|
|
)
|
|
def test_hardware_compatibility_level(self, identity_builder_network, level):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(hardware_compatibility_level=level)
|
|
with loader(builder, network) as config:
|
|
assert config.hardware_compatibility_level == level
|
|
|
|
if mod.version(trt.__version__) >= mod.version("10.2"):
|
|
|
|
@pytest.mark.parametrize(
|
|
"platform",
|
|
[
|
|
trt.RuntimePlatform.SAME_AS_BUILD,
|
|
trt.RuntimePlatform.WINDOWS_AMD64,
|
|
],
|
|
)
|
|
def test_runtime_platform(self, identity_builder_network, platform):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(runtime_platform=platform)
|
|
with loader(builder, network) as config:
|
|
assert config.runtime_platform == platform
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="Unsupported for TRT versions prior to 8.6",
|
|
)
|
|
@pytest.mark.parametrize("num_streams", range(3))
|
|
def test_max_aux_streams(self, identity_builder_network, num_streams):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(max_aux_streams=num_streams)
|
|
with loader(builder, network) as config:
|
|
assert config.max_aux_streams == num_streams
|
|
|
|
@pytest.mark.skipif(
|
|
mod.version(trt.__version__) < mod.version("9.0") and not polygraphy_config.USE_TENSORRT_RTX,
|
|
reason="API was added in TRT 9.0",
|
|
)
|
|
def test_progress_monitor(self, identity_builder_network):
|
|
class DummyProgressMonitor(trt.IProgressMonitor):
|
|
def __init__(self):
|
|
trt.IProgressMonitor.__init__(self)
|
|
|
|
def phase_start(self, phase_name, parent_phase, num_steps):
|
|
pass
|
|
|
|
def phase_finish(self, phase_name):
|
|
pass
|
|
|
|
def step_complete(self, phase_name, step):
|
|
return True
|
|
|
|
builder, network = identity_builder_network
|
|
progress_monitor = DummyProgressMonitor()
|
|
loader = CreateConfig(progress_monitor=progress_monitor)
|
|
with loader(builder, network) as config:
|
|
assert config.progress_monitor == progress_monitor
|
|
|
|
if mod.version(trt.__version__) >= mod.version("10.8") and not polygraphy_config.USE_TENSORRT_RTX:
|
|
@pytest.mark.parametrize(
|
|
"level",
|
|
[
|
|
trt.TilingOptimizationLevel.NONE,
|
|
trt.TilingOptimizationLevel.FAST,
|
|
trt.TilingOptimizationLevel.MODERATE,
|
|
trt.TilingOptimizationLevel.FULL,
|
|
],
|
|
)
|
|
def test_tiling_optimization_level(self, identity_builder_network, level):
|
|
builder, network = identity_builder_network
|
|
loader = CreateConfig(tiling_optimization_level=level)
|
|
with loader(builder, network) as config:
|
|
assert config.tiling_optimization_level == level
|
|
|
|
|
|
class TestPostprocessConfig:
|
|
def test_with_config(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
config = CreateConfig()(builder, network)
|
|
assert not config.get_flag(trt.BuilderFlag.INT8)
|
|
|
|
config = postprocess_config(
|
|
config,
|
|
func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
|
|
builder=builder,
|
|
network=network,
|
|
)
|
|
assert config.get_flag(trt.BuilderFlag.INT8)
|
|
|
|
def test_with_config_callable(self, identity_builder_network):
|
|
builder, network = identity_builder_network
|
|
config = CreateConfig()
|
|
|
|
config = postprocess_config(
|
|
config,
|
|
func=lambda builder, network, config: config.set_flag(trt.BuilderFlag.INT8),
|
|
builder=builder,
|
|
network=network,
|
|
)
|
|
assert config.get_flag(trt.BuilderFlag.INT8)
|