wehub-resource-sync
126 lines
4.2 KiB
Python
126 lines
4.2 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 subprocess as sp
|
|
|
|
import numpy as np
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
import onnxruntime
|
|
import pytest
|
|
from onnx_graphsurgeon.logger import G_LOGGER
|
|
from onnx_graphsurgeon.util import misc
|
|
|
|
ROOT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), os.path.pardir))
|
|
EXAMPLES_ROOT = os.path.join(ROOT_DIR, "examples")
|
|
|
|
|
|
class Artifact(object):
|
|
def __init__(self, name, infer=True):
|
|
self.name = name
|
|
self.infer = infer
|
|
|
|
|
|
EXAMPLES = [
|
|
("01_creating_a_model", [Artifact("test_globallppool.onnx")]),
|
|
("02_creating_a_model_with_initializer", [Artifact("test_conv.onnx")]),
|
|
("03_isolating_a_subgraph", [Artifact("model.onnx"), Artifact("subgraph.onnx")]),
|
|
("04_modifying_a_model", [Artifact("model.onnx"), Artifact("modified.onnx")]),
|
|
("05_folding_constants", [Artifact("model.onnx"), Artifact("folded.onnx")]),
|
|
(
|
|
"06_removing_nodes",
|
|
[Artifact("model.onnx", infer=False), Artifact("removed.onnx")],
|
|
),
|
|
("07_creating_a_model_with_the_layer_api", [Artifact("model.onnx")]),
|
|
("08_replacing_a_subgraph", [Artifact("model.onnx"), Artifact("replaced.onnx")]),
|
|
("09_shape_operations_with_the_layer_api", [Artifact("model.onnx")]),
|
|
("10_dynamic_batch_size", [Artifact("model.onnx"), Artifact("dynamic.onnx")]),
|
|
("11_creating_a_local_function", [Artifact("model.onnx")]),
|
|
# Skipping inference test as bf16 is not supported in ORT yet.
|
|
(
|
|
"12_using_numpy_unsupported_dtypes",
|
|
[Artifact("test_conv_bf16.onnx", infer=False)],
|
|
),
|
|
]
|
|
|
|
|
|
# Extract any ``` blocks from the README
|
|
def load_commands_from_readme(readme):
|
|
def ignore_command(cmd):
|
|
return "pip" in cmd
|
|
|
|
commands = []
|
|
with open(readme, "r") as f:
|
|
in_command_block = False
|
|
for line in f.readlines():
|
|
if not in_command_block and "```bash" in line:
|
|
in_command_block = True
|
|
elif in_command_block:
|
|
if "```" in line:
|
|
in_command_block = False
|
|
elif not ignore_command(line):
|
|
commands.append(line.strip())
|
|
return commands
|
|
|
|
|
|
def infer_model(path):
|
|
model = onnx.load(path)
|
|
onnx.checker.check_model(model)
|
|
|
|
graph = gs.import_onnx(model)
|
|
|
|
feed_dict = {}
|
|
for tensor in graph.inputs:
|
|
shape = tuple(
|
|
dim if not misc.is_dynamic_dimension(dim) else 1 for dim in tensor.shape
|
|
)
|
|
feed_dict[tensor.name] = np.random.random_sample(size=shape).astype(
|
|
tensor.dtype
|
|
)
|
|
|
|
output_names = [out.name for out in graph.outputs]
|
|
|
|
sess = onnxruntime.InferenceSession(
|
|
model.SerializeToString(), providers=["CPUExecutionProvider"]
|
|
)
|
|
outputs = sess.run(output_names, feed_dict)
|
|
G_LOGGER.info("Inference outputs: {:}".format(outputs))
|
|
return outputs
|
|
|
|
|
|
@pytest.mark.parametrize("example_dir,artifacts", EXAMPLES)
|
|
def test_examples(example_dir, artifacts):
|
|
example_dir = os.path.join(EXAMPLES_ROOT, example_dir)
|
|
readme = os.path.join(example_dir, "README.md")
|
|
commands = load_commands_from_readme(readme)
|
|
for command in commands:
|
|
G_LOGGER.info(command)
|
|
assert (
|
|
sp.run(
|
|
["bash", "-c", command], cwd=example_dir, env={"PYTHONPATH": ROOT_DIR}
|
|
).returncode
|
|
== 0
|
|
)
|
|
|
|
for artifact in artifacts:
|
|
artifact_path = os.path.join(example_dir, artifact.name)
|
|
assert os.path.exists(artifact_path)
|
|
if artifact.infer:
|
|
assert infer_model(artifact_path)
|
|
os.remove(artifact_path)
|