"""Standalone synthetic TensorRT If -> Transpose -> Resize reproducer."""
import argparse
import ctypes
import ctypes.util
from pathlib import Path
import numpy as np
import onnx
from onnx import helper as h
from onnx import numpy_helper as nh
def model(batch: int, mode: str, transpose: bool = True, before: bool = False) -> onnx.ModelProto:
vi = h.make_tensor_value_info
branch_shape = [batch, 1, "W", "H"] if transpose and before else [batch, 1, "H", "W"]
branch_input = "ORIENTED" if transpose and before else "X"
yes = h.make_graph([h.make_node("Add", [branch_input, "TWO"], ["YES"])], "then", [], [vi("YES", 1, branch_shape)])
no = h.make_graph([h.make_node("Add", [branch_input, "ONE"], ["NO"])], "else", [], [vi("NO", 1, branch_shape)])
nodes = [h.make_node("If", ["COND"], ["SELECTED"], then_branch=yes, else_branch=no)]
if transpose and before:
nodes.insert(0, h.make_node("Transpose", ["X"], ["ORIENTED"], perm=[0, 1, 3, 2]))
if transpose and not before:
nodes.append(h.make_node("Transpose", ["SELECTED"], ["ORIENTED"], perm=[0, 1, 3, 2]))
nodes.extend(
[
h.make_node("Cast", ["SIZE"], ["SIZE64"], to=7),
h.make_node("Concat", ["NC", "SIZE64"], ["TARGET"], axis=0),
h.make_node(
"Resize",
["ORIENTED" if transpose and not before else "SELECTED", "", "", "TARGET"],
["Y"],
mode=mode,
coordinate_transformation_mode="align_corners",
cubic_coeff_a=-0.75,
),
]
)
weights = [
nh.from_array(np.full((1, 1, 1, 1), 2.0, np.float32), "TWO"),
nh.from_array(np.ones((1, 1, 1, 1), np.float32), "ONE"),
nh.from_array(np.array([batch, 1], np.int64), "NC"),
]
result = h.make_model(
h.make_graph(
nodes,
"conditional-transpose",
[vi("X", 1, [batch, 1, "H", "W"]), vi("COND", 9, []), vi("SIZE", 6, [2])],
[vi("Y", 1, [batch, 1, "OH", "OW"])],
initializer=weights,
),
opset_imports=[h.make_opsetid("", 18)],
ir_version=10,
)
onnx.checker.check_model(result, full_check=True)
return result
def build(args: argparse.Namespace) -> None:
import tensorrt as trt # noqa: PLC0415 - optional GPU dependency
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(0)
if args.api:
x = network.add_input("X", trt.float32, (args.batch, 1, -1, -1))
cond = network.add_input("COND", trt.bool, ())
size = network.add_input("SIZE", trt.int32, (2,))
if not args.no_transpose and args.before:
orient = network.add_shuffle(x)
orient.second_transpose = (0, 1, 3, 2)
x = orient.get_output(0)
branch = network.add_if_conditional()
branch.set_condition(cond)
inside = branch.add_input(x).get_output(0)
one = np.ones((1, 1, 1, 1), np.float32)
two = one * 2
yes = network.add_elementwise(
inside, network.add_constant(two.shape, two).get_output(0), trt.ElementWiseOperation.SUM
)
no = network.add_elementwise(
inside, network.add_constant(one.shape, one).get_output(0), trt.ElementWiseOperation.SUM
)
value = branch.add_output(yes.get_output(0), no.get_output(0)).get_output(0)
if not args.no_transpose and not args.before:
orient = network.add_shuffle(value)
orient.second_transpose = (0, 1, 3, 2)
value = orient.get_output(0)
nc = np.array([args.batch, 1], np.int32)
sizes = network.add_concatenation([network.add_constant(nc.shape, nc).get_output(0), size])
sizes.axis = 0
resize = network.add_resize(value)
resize.resize_mode = getattr(trt.InterpolationMode, args.interpolation.upper())
resize.coordinate_transformation = trt.ResizeCoordinateTransformation.ALIGN_CORNERS
resize.cubic_coeff = -0.75
resize.set_input(1, sizes.get_output(0))
result = resize.get_output(0)
result.name = "Y"
network.mark_output(result)
else:
parser = trt.OnnxParser(network, logger)
assert parser.parse(args.onnx.read_bytes())
config = builder.create_builder_config()
config.clear_flag(trt.BuilderFlag.TF32)
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 512 << 20)
config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
profile = builder.create_optimization_profile()
profile.set_shape(
"X",
(args.batch, 1, 1, 1),
(args.batch, 1, args.opt_height, args.opt_width),
(args.batch, 1, 1536, 3072),
)
profile.set_shape_input("SIZE", (1, 1), (512, 512), (3072, 3072))
config.add_optimization_profile(profile)
plan = builder.build_serialized_network(network, config)
assert plan is not None
args.plan.write_bytes(bytes(plan))
def run(a: argparse.Namespace) -> None: # noqa: PLR0915 - keep raw CUDA execution explicit
import onnxruntime as ort # noqa: PLC0415 - only required by the runner
import tensorrt as trt # noqa: PLC0415 - optional GPU dependency
cuda = ctypes.CDLL(a.cudart or ctypes.util.find_library("cudart") or "libcudart.so")
for name, args in [
("cudaSetDevice", [ctypes.c_int]),
("cudaMalloc", [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]),
("cudaFree", [ctypes.c_void_p]),
("cudaMemcpy", [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]),
("cudaStreamCreate", [ctypes.POINTER(ctypes.c_void_p)]),
("cudaStreamSynchronize", [ctypes.c_void_p]),
("cudaDeviceSynchronize", []),
("cudaStreamDestroy", [ctypes.c_void_p]),
]:
fn = getattr(cuda, name)
fn.argtypes = args
fn.restype = ctypes.c_int
def call(name: str, *args: object) -> None:
code = getattr(cuda, name)(*args)
if code:
raise RuntimeError(f"{name}: CUDA {code}")
def allocate(array: np.ndarray) -> ctypes.c_void_p:
ptr = ctypes.c_void_p()
call("cudaMalloc", ctypes.byref(ptr), array.nbytes)
call("cudaMemcpy", ptr, ctypes.c_void_p(array.ctypes.data), array.nbytes, 1)
return ptr
call("cudaSetDevice", 0)
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(a.plan.read_bytes())
context = engine.create_execution_context()
stream = ctypes.c_void_p()
if not a.default_stream:
call("cudaStreamCreate", ctypes.byref(stream))
handle = stream.value or 0
assert context.set_optimization_profile_async(0, handle)
call("cudaStreamSynchronize", stream)
opt = ort.SessionOptions()
opt.intra_op_num_threads = 1
ref = ort.InferenceSession(str(a.onnx), sess_options=opt, providers=["CPUExecutionProvider"])
x = (np.arange(a.batch * a.height * a.width, dtype=np.float32).reshape(a.batch, 1, a.height, a.width) % 113) * 0.1
for condition in (True, False, True):
feed = {"X": x, "COND": np.array(condition, np.bool_), "SIZE": np.array([19, 33], np.int32)}
expected = ref.run(None, feed)[0]
allocs = {}
for name, value in feed.items():
ptr = ctypes.c_void_p(value.ctypes.data)
if engine.get_tensor_location(name) == trt.TensorLocation.DEVICE:
ptr = allocate(value)
allocs[name] = ptr
assert context.set_input_shape(name, tuple(value.shape))
assert context.set_tensor_address(name, ptr.value)
assert not context.infer_shapes()
actual = np.full(tuple(context.get_tensor_shape("Y")), np.nan, np.float32)
out = allocate(actual)
allocs["Y"] = out
assert context.set_tensor_address("Y", out.value)
call("cudaDeviceSynchronize")
print("execute raw", a.batch, a.api, condition, flush=True)
assert context.execute_async_v3(handle)
call("cudaStreamSynchronize", stream)
call("cudaMemcpy", ctypes.c_void_p(actual.ctypes.data), out, actual.nbytes, 2)
print(
"actual",
actual.min(),
actual.max(),
"expected",
expected.min(),
expected.max(),
"max_error",
np.max(abs(actual - expected)),
flush=True,
)
for ptr in allocs.values():
call("cudaFree", ptr)
np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-4)
print("PASS raw", flush=True)
if not a.default_stream:
call("cudaStreamDestroy", stream)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("action", choices=("build", "run"))
parser.add_argument("--batch", type=int, default=4)
parser.add_argument("--height", type=int, default=19)
parser.add_argument("--width", type=int, default=31)
parser.add_argument("--opt-height", type=int, default=1024)
parser.add_argument("--opt-width", type=int, default=1024)
parser.add_argument("--api", action="store_true", help="Bypass the ONNX parser when building TensorRT")
parser.add_argument("--no-transpose", action="store_true")
parser.add_argument("--before", action="store_true", help="Move Transpose before If (control case)")
parser.add_argument("--default-stream", action="store_true")
parser.add_argument("--cudart", help="Path to libcudart.so when unavailable on the library search path")
parser.add_argument("--interpolation", choices=("cubic", "linear"), default="cubic")
parser.add_argument("--directory", type=Path, default=Path("dist/source-prepost/conditional-investigation/repro"))
args = parser.parse_args()
args.directory.mkdir(parents=True, exist_ok=True)
order = "plain" if args.no_transpose else ("transpose-before" if args.before else "transpose")
stem = f"b{args.batch}-{args.interpolation}-{order}"
args.onnx = args.directory / f"{stem}.onnx"
args.plan = args.directory / f"{stem}-{'api' if args.api else 'onnx'}.plan"
if args.action == "build":
onnx.save(model(args.batch, args.interpolation, not args.no_transpose, args.before), args.onnx)
build(args)
else:
run(args)
if __name__ == "__main__":
main()
Description
TensorRT successfully builds a small FP32
If → Transpose → Resizenetwork, but produces substantially incorrect results at a valid dynamic input shape. This reproduces with 11.1.0.106 on A10G and GTX 1660 SUPER, and persists with 11.2.1.2 on GTX 1660 SUPER, through both the ONNX parser and the direct TensorRT network API. A10G has not been tested with 11.2.The operation is equivalent to:
At runtime shape
[4,1,19,31], the maximum absolute error against ONNX Runtime CPU is 14.6497345. Moving Transpose before If makes the same comparison pass.Environment
tensorrt-cu13)Minimal reproducer
Save the following code as
reproduce_trt_conditional_transpose.py. It generates synthetic inputs, the ONNX reference and the TensorRT engine. With--api, the TensorRT build uses the network API directly and does not invoke the ONNX parser.Standalone Python reproducer
Steps to reproduce
Install the dependencies in a Python 3.13 environment on a CUDA-capable Linux host:
For the 11.1 comparison, use
tensorrt-cu13==11.1.0.106in a separate environment. Locate the installed CUDA runtime and build/run as separate processes:An existing system
libcudart.so.13path can be supplied instead. The script also supports automatic library lookup when--cudartis omitted.The default case uses these settings:
[4,1,H,W][2][4,1,19,31]/[19,33](arange(4*19*31).reshape(4,1,19,31) % 113) * 0.1, FP32Omit
--apifrom both commands to reproduce through the ONNX parser. Use a separate directory when changing profile dimensions, since the filename does not encode the profile.Expected behavior
TensorRT should match the ONNX Runtime CPU result within
rtol=1e-5, atol=1e-4for this valid input shape. The reference output for the first True condition has range 0.5502561…14.6497345.Actual behavior
Engine construction succeeds, but the first inference produces range
−0.38888925…3.607245, with maximum absolute error 14.6497345:
The B4 direct-API and ONNX paths produce these same values on both tested GPUs. A second fresh-process execution of the B4 direct-API engine on A10G reproduces the same failure.
Have you tried the latest release?: Yes
Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (
polygraphy run <model.onnx> --onnxrt): Yes