466 lines
15 KiB
Python
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,
|
|
),
|
|
}
|