Files
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

466 lines
15 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 os
import numpy as np
import tensorrt as trt
from polygraphy import mod, util
from polygraphy.backend.common import BytesFromPath
from polygraphy.backend.onnx import OnnxFromPath
from polygraphy.backend.tf import GraphFromFrozen
from polygraphy.common import TensorMetadata
from polygraphy.datatype import DataType
def model_path(name=None):
path = os.path.abspath(os.path.dirname(__file__))
if name is not None:
path = os.path.join(path, name)
return path
class Model:
def __init__(
self, path, LoaderType, check_runner, input_metadata=None, ext_data=None
):
self.path = path
self.loader = LoaderType(self.path)
self.check_runner = check_runner
self.input_metadata = input_metadata
self.ext_data = ext_data
def check_tf_identity(runner):
feed_dict = {
"Input:0": np.random.random_sample(size=(1, 15, 25, 30)).astype(np.float32)
}
outputs = runner.infer(feed_dict)
assert np.all(outputs["Identity_2:0"] == feed_dict["Input:0"])
MODELS_DIR = os.path.join(os.path.dirname(__file__))
TF_MODELS = {
"identity": Model(
path=model_path("tf_identity.pb"),
LoaderType=GraphFromFrozen,
check_runner=check_tf_identity,
),
}
def check_identity(runner):
feed_dict = {"x": np.random.random_sample(size=(1, 1, 2, 2)).astype(np.float32)}
outputs = runner.infer(feed_dict)
assert np.all(outputs["y"] == feed_dict["x"])
def check_identity_identity(runner):
feed_dict = {"X": np.random.random_sample(size=(64, 64)).astype(np.float32)}
outputs = runner.infer(feed_dict)
assert np.all(outputs["identity_out_2"] == feed_dict["X"])
def check_dynamic_identity(runner, shapes):
feed_dict = {"X": np.random.random_sample(size=shapes["X"]).astype(np.float32)}
outputs = runner.infer(feed_dict)
assert np.array_equal(outputs["Y"], feed_dict["X"])
def check_empty_tensor_expand(runner, shapes):
shape = shapes["new_shape"]
feed_dict = {
"data": np.zeros(shape=(2, 0, 3, 0), dtype=np.float32),
"new_shape": np.array(
shape,
dtype=(
np.int32
if mod.version(trt.__version__) < mod.version("9.0")
else np.int64
),
),
}
outputs = runner.infer(feed_dict)
# Empty tensor will still be empty after broadcast
assert outputs["expanded"].shape == shape
assert util.volume(outputs["expanded"].shape) == 0
def check_reshape(runner):
feed_dict = {"data": np.random.random_sample(size=(1, 3, 5, 5)).astype(np.float32)}
outputs = runner.infer(feed_dict)
assert np.all(outputs["output"] == feed_dict["data"].ravel())
def check_residual_block(runner, shapes):
feed_dict = {
"gpu_0/data_0": np.random.random_sample(size=shapes["gpu_0/data_0"]).astype(
np.float32
)
}
# Confirm inference can go through without error
outputs = runner.infer(feed_dict)
def check_matmul_2layer(runner, shape=(2, 8)):
feed_dict = {
"onnx::MatMul_0": np.random.random_sample(size=shape).astype(np.float32)
}
# Confirm inference can go through without error
outputs = runner.infer(feed_dict)
def no_check_implemented(runner):
raise NotImplementedError("No check_runner implemented for this model")
ONNX_MODELS = {
"identity": Model(
path=model_path("identity.onnx"),
LoaderType=BytesFromPath,
check_runner=check_identity,
input_metadata=TensorMetadata().add(
"x", dtype=DataType.FLOAT32, shape=(1, 1, 2, 2)
),
),
"identity_identity": Model(
path=model_path("identity_identity.onnx"),
LoaderType=BytesFromPath,
check_runner=check_identity_identity,
),
"dynamic_identity": Model(
path=model_path("dynamic_identity.onnx"),
LoaderType=BytesFromPath,
check_runner=check_dynamic_identity,
input_metadata=TensorMetadata().add(
"X", dtype=DataType.FLOAT32, shape=(1, 1, -1, -1)
),
),
"identity_multi_ch": Model(
path=model_path("identity_multi_ch.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
input_metadata=TensorMetadata().add(
"x", dtype=DataType.FLOAT32, shape=(2, 4, 3, 3)
),
),
"empty_tensor_expand": Model(
path=model_path("empty_tensor_expand.onnx"),
LoaderType=BytesFromPath,
check_runner=check_empty_tensor_expand,
),
"and": Model(
path=model_path("and.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"scan": Model(
path=model_path("scan.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"pow_scalar": Model(
path=model_path("pow_scalar.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"dim_param": Model(
path=model_path("dim_param.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"tensor_attr": Model(
path=model_path("tensor_attr.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"identity_with_initializer": Model(
path=model_path("identity_with_initializer.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"const_foldable": Model(
path=model_path("const_foldable.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"reshape": Model(
path=model_path("reshape.onnx"),
LoaderType=BytesFromPath,
check_runner=check_reshape,
),
"reducable": Model(
path=model_path("reducable.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
input_metadata=TensorMetadata()
.add("X0", shape=(1,), dtype=DataType.FLOAT32)
.add("Y0", shape=(1,), dtype=DataType.FLOAT32),
),
"reducable_with_const": Model(
path=model_path("reducable_with_const.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"ext_weights": Model(
path=model_path("ext_weights.onnx"),
LoaderType=OnnxFromPath,
check_runner=no_check_implemented,
ext_data=model_path("data"),
),
"ext_weights_same_dir": Model(
path=model_path(os.path.join("ext_weights_same_dir", "ext_weights.onnx")),
LoaderType=OnnxFromPath,
check_runner=no_check_implemented,
ext_data=model_path("ext_weights_same_dir"),
),
"capability": Model(
path=model_path("capability.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"instancenorm": Model(
path=model_path("instancenorm.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"add_with_dup_inputs": Model(
path=model_path("add_with_dup_inputs.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"needs_constraints": Model(
path=model_path("needs_constraints.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
input_metadata=TensorMetadata().add(
"x", dtype=DataType.FLOAT32, shape=(1, 1, 256, 256)
),
),
"constant_fold_bloater": Model(
path=model_path("constant_fold_bloater.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"renamable": Model(
path=model_path("renamable.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"cleanable": Model(
path=model_path("cleanable.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"nonzero": Model(
path=model_path("nonzero.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"inp_dim_val_not_set": Model(
path=model_path("inp_dim_val_not_set.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"multi_output": Model(
path=model_path("multi_output.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"unbounded_dds": Model(
path=model_path("unbounded_dds.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"loop": Model(
path=model_path("loop.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"matmul.fp16": Model(
path=model_path("matmul.fp16.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"matmul": Model(
path=model_path("matmul.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"sparse.matmul": Model(
path=model_path("sparse.matmul.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"matmul.bf16": Model(
path=model_path("matmul.bf16.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"matmul.bf16.i32data": Model(
path=model_path("matmul.bf16.i32data.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"matmul_2layer": Model(
path=model_path("matmul_2layer.onnx"),
LoaderType=BytesFromPath,
check_runner=check_matmul_2layer,
),
"unsorted": Model(
path=model_path("unsorted.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"conv": Model(
path=model_path("conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"sparse.conv": Model(
path=model_path("sparse.conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"no_op_reshape": Model(
path=model_path("no_op_reshape.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_dup_value_info": Model(
path=model_path("bad_graph_with_dup_value_info.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_no_name": Model(
path=model_path("bad_graph_with_no_name.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_no_import_domains": Model(
path=model_path("bad_graph_with_no_import_domains.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_parallel_invalid_nodes": Model(
path=model_path("bad_graph_with_parallel_invalid_nodes.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_conditionally_invalid": Model(
path=model_path("bad_graph_conditionally_invalid.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"custom_op_node": Model(
path=model_path("custom_op_node.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_duplicate_node_names": Model(
path=model_path("bad_graph_with_duplicate_node_names.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"bad_graph_with_multi_level_errors": Model(
path=model_path("bad_graph_with_multi_level_errors.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"empty": Model(
path=model_path("empty.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"residual_block": Model(
path=model_path("residual_block.onnx"),
LoaderType=BytesFromPath,
check_runner=check_residual_block,
),
"graph_with_subgraph_matching_toy_plugin": Model(
path=model_path("toy_subgraph.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"transpose_matmul": Model(
path=model_path("transpose_matmul.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"qdq_conv": Model(
path=model_path("qdq_conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.matmul.fp16": Model(
path=model_path("weightless.matmul.fp16.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.matmul.bf16": Model(
path=model_path("weightless.matmul.bf16.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.conv": Model(
path=model_path("weightless.conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.sparse.matmul": Model(
path=model_path("weightless.sparse.matmul.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.sparse.conv": Model(
path=model_path("weightless.sparse.conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.transpose_matmul": Model(
path=model_path("weightless.transpose_matmul.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"weightless.qdq_conv": Model(
path=model_path("weightless.qdq_conv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"roialign": Model(
path=model_path("roialign.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"attention": Model(
path=model_path("attention.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"multi_attention": Model(
path=model_path("multi_attention.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
"attention_same_qkv": Model(
path=model_path("attention_same_qkv.onnx"),
LoaderType=BytesFromPath,
check_runner=no_check_implemented,
),
}