nvidia--tensorrt/tools/Polygraphy/tests/backend/onnx/test_loader.py

#
# 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 tempfile

import numpy as np
import onnx
import onnx_graphsurgeon as gs
import pytest

from polygraphy import constants
from polygraphy.backend.onnx import (
    ConvertToFp16,
    FoldConstants,
    ModifyOutputs,
    OnnxFromBytes,
    OnnxFromPath,
    OnnxFromTfGraph,
    SaveOnnx,
    SetUpperBound,
    extract_subgraph,
    fold_constants,
    gs_from_onnx,
    infer_shapes,
    onnx_from_path,
)
from polygraphy.common import TensorMetadata
from polygraphy.logger import G_LOGGER
from tests.helper import is_file_non_empty
from tests.models.meta import ONNX_MODELS, TF_MODELS


class TestLoggerCallbacks:
    @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys())
    def test_set_severity(self, sev):
        G_LOGGER.module_severity = sev


class TestOnnxFromPath:
    def test_basic(self):
        loader = OnnxFromPath(ONNX_MODELS["identity"].path)
        model = loader()
        assert isinstance(model, onnx.ModelProto)
        assert len(model.graph.node) == 1

    @pytest.mark.serial
    def test_warn_if_impl_methods_called(self, check_warnings_on_loader_impl_methods):
        check_warnings_on_loader_impl_methods(
            OnnxFromPath(ONNX_MODELS["identity"].path)
        )

    def test_external_data(self):
        model = ONNX_MODELS["ext_weights"]
        loader = OnnxFromPath(model.path, model.ext_data)
        assert isinstance(loader(), onnx.ModelProto)

    def test_ignore_external_data(self):
        model = ONNX_MODELS["ext_weights"]
        loader = OnnxFromPath(model.path, ignore_external_data=True)
        onnx_model = loader()

        assert isinstance(onnx_model, onnx.ModelProto)
        assert all(init.data_location == 1 for init in onnx_model.graph.initializer)


class TestOnnxFromBytes:
    def test_basic(self):
        loader = OnnxFromBytes(ONNX_MODELS["identity"].loader)
        model = loader()
        assert isinstance(model, onnx.ModelProto)
        assert len(model.graph.node) == 1


class TestGsFromOnnx:
    def test_basic(self):
        graph = gs_from_onnx(OnnxFromPath(ONNX_MODELS["identity"].path))
        assert isinstance(graph, gs.Graph)


class TestExportOnnxFromTf:
    def test_no_optimize(self):
        pytest.importorskip("tensorflow")
        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, optimize=False)
        model = loader()

    def test_opset(self):
        pytest.importorskip("tensorflow")
        loader = OnnxFromTfGraph(TF_MODELS["identity"].loader, opset=9)
        model = loader()
        assert model.opset_import[0].version == 9


class TestModifyOnnx:
    @pytest.mark.parametrize("copy", [True, False])
    def test_layerwise(self, copy):
        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
        loader = ModifyOutputs(original_model, outputs=constants.MARK_ALL, copy=copy)
        model = loader()
        assert len(original_model.graph.output) == 1 or not copy
        assert len(model.graph.output) == 2

    @pytest.mark.parametrize("output", ["identity_out_0", "identity_out_2"])
    def test_custom_outputs(self, output):
        loader = ModifyOutputs(
            OnnxFromPath(ONNX_MODELS["identity_identity"].path), outputs=[output]
        )
        model = loader()
        assert len(model.graph.output) == 1
        assert model.graph.output[0].name == output

    def test_exclude_outputs_with_layerwise(self):
        loader = ModifyOutputs(
            OnnxFromPath(ONNX_MODELS["identity_identity"].path),
            outputs=constants.MARK_ALL,
            exclude_outputs=["identity_out_2"],
        )
        model = loader()
        assert len(model.graph.output) == 1
        assert model.graph.output[0].name == "identity_out_0"


@pytest.mark.parametrize("allow_onnxruntime", [True, False])
class TestInferShapes:
    def check_model(self, model):
        # Find all intermediate tensors to check if they have shapes.
        tensors = set()
        for node in model.graph.node:
            tensors.update(node.output)
        tensors -= {out.name for out in model.graph.output}

        assert len(model.graph.value_info) >= len(tensors)
        for val in model.graph.value_info:
            assert val.type.tensor_type.HasField("shape")

    def test_model(self, allow_onnxruntime):
        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
        model = infer_shapes(original_model, allow_onnxruntime=allow_onnxruntime)
        self.check_model(model)

    def test_path(self, allow_onnxruntime):
        model = infer_shapes(
            ONNX_MODELS["identity_identity"].path, allow_onnxruntime=allow_onnxruntime
        )
        self.check_model(model)

    @pytest.mark.parametrize("set_data_dir", [True, False])
    def test_external_data(self, set_data_dir, allow_onnxruntime):
        model = ONNX_MODELS["ext_weights_same_dir"]
        model = infer_shapes(
            model.path,
            external_data_dir=model.ext_data if set_data_dir else None,
            allow_onnxruntime=allow_onnxruntime,
        )
        self.check_model(model)

    def test_save_to_disk_on_size_threshold(self, allow_onnxruntime):
        model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
        model = infer_shapes(
            model, save_to_disk_threshold_bytes=0, allow_onnxruntime=allow_onnxruntime
        )
        self.check_model(model)


class TestConvertToFp16:
    @pytest.mark.parametrize("copy", [True, False])
    def test_basic(self, copy):
        # Precondition.
        original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path)
        assert original_model.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT or not copy

        # Under test.
        loader = ConvertToFp16(original_model, copy=copy)
        model = loader()

        # Postcondition.
        graph = gs_from_onnx(model)
        graph.toposort()

        assert graph.inputs[0].dtype == "float32"
        assert graph.nodes[0].op == "Cast"
        assert graph.nodes[1].op == "Identity"
        assert graph.nodes[2].op == "Identity"
        assert graph.nodes[3].op == "Cast"
        assert graph.outputs[0].dtype == "float32"


class TestFoldConstants:
    @pytest.mark.parametrize("fold_shapes", [True, False])
    @pytest.mark.parametrize("partitioning", [None, "basic", "recursive"])
    @pytest.mark.parametrize("copy", [True, False])
    @pytest.mark.parametrize("allow_onnxruntime_shape_inference", [True, False])
    def test_basic(
        self, partitioning, fold_shapes, copy, allow_onnxruntime_shape_inference
    ):
        original_model = onnx_from_path(ONNX_MODELS["const_foldable"].path)
        loader = FoldConstants(
            original_model,
            partitioning=partitioning,
            fold_shapes=fold_shapes,
            copy=copy,
            error_ok=False,
            allow_onnxruntime_shape_inference=allow_onnxruntime_shape_inference,
        )
        model = loader()
        assert len(original_model.graph.node) != 1 or not copy
        assert len(model.graph.node) == 1

    @pytest.mark.parametrize(
        "size_threshold, expect_folding",
        [
            (None, True),
            (0, False),
            (10 << 20, True),
            (10 << 20 - 1, False),
        ],
    )
    def test_size_threshold(self, size_threshold, expect_folding):
        model = onnx_from_path(ONNX_MODELS["constant_fold_bloater"].path)
        model = fold_constants(model, size_threshold=size_threshold)

        if expect_folding:
            assert len(model.graph.node) == 0
        else:
            assert len(model.graph.node) == 1
            assert model.graph.node[0].op_type == "Tile"


class TestSetUpperBound:
    @pytest.mark.parametrize("global_upper_bound", [False, True])
    @pytest.mark.parametrize("specified_upper_bound", [False, True])
    def test_set_upper_bound(
        self,
        global_upper_bound,
        specified_upper_bound,
    ):
        original_model = onnx_from_path(ONNX_MODELS["unbounded_dds"].path)
        upper_bound_dict = {}
        if not global_upper_bound and not specified_upper_bound:
            upper_bound_dict[""] = 1000
            upper_bound = 1000
        if global_upper_bound:
            upper_bound_dict[""] = 2000
            upper_bound = 2000
        if specified_upper_bound:
            upper_bound_dict["cast_out_6"] = 4000
            upper_bound = 4000

        loader = SetUpperBound(
            original_model,
            upper_bounds=upper_bound_dict,
        )

        model = loader()
        graph = gs_from_onnx(model)

        # Check if there is a Min operator in the modified model
        find_min = False
        for node in graph.nodes:
            if node.op == "Min":
                find_min = True
                # Check if the Min operator's second input is a constant tensor
                assert isinstance(node.inputs[1], gs.Constant)

                val = node.inputs[1].values
                # Check if the constant value equals the target upper bound
                assert val == upper_bound
        assert find_min


class TestSaveOnnx:
    def test_save_onnx(self):
        with tempfile.TemporaryDirectory() as outdir:
            outpath = os.path.join(outdir, "test", "nested")
            loader = SaveOnnx(OnnxFromPath(ONNX_MODELS["identity"].path), path=outpath)
            loader()
            assert is_file_non_empty(outpath)

    def test_external_data(self):
        with tempfile.NamedTemporaryFile(dir=".") as path, tempfile.NamedTemporaryFile(dir=".") as data:
            rpath_name = os.path.basename(data.name)
            model = OnnxFromPath(ONNX_MODELS["const_foldable"].path)
            loader = SaveOnnx(
                model, path.name, external_data_path=rpath_name, size_threshold=0
            )
            loader()
            assert is_file_non_empty(path.name)
            assert is_file_non_empty(data.name)


@pytest.fixture()
def extract_model():
    input_metadata = TensorMetadata().add("X", dtype=np.float32, shape=(64, 64))
    output_metadata = TensorMetadata().add(
        "identity_out_0", dtype=np.float32, shape=None
    )
    return (
        onnx_from_path(ONNX_MODELS["identity_identity"].path),
        input_metadata,
        output_metadata,
    )


class TestExtractSubgraph:
    def check_model(self, model):
        graph = gs_from_onnx(model)
        assert len(graph.nodes) == 1

        assert len(graph.inputs) == 1
        assert graph.inputs[0].name == "X"
        assert graph.inputs[0].shape is not None
        assert graph.inputs[0].dtype is not None

        assert len(graph.outputs) == 1
        assert graph.outputs[0].name == "identity_out_0"
        assert graph.outputs[0].dtype is not None

    def test_extract_onnx_model(self, extract_model):
        original_model, input_meta, output_meta = extract_model
        model = extract_subgraph(original_model, input_meta, output_meta)

        assert original_model.graph.output[0].name == "identity_out_2"
        self.check_model(model)

    def test_extract_onnx_model_no_input_meta(self, extract_model):
        model, _, output_meta = extract_model
        model = extract_subgraph(model, output_metadata=output_meta)
        self.check_model(model)

    def test_extract_onnx_model_no_output_meta(self, extract_model):
        model, input_meta, _ = extract_model
        model = extract_subgraph(model, input_metadata=input_meta)
        assert model.graph.output[0].name == "identity_out_2"

    def test_extract_onnx_gs_graph(self, extract_model):
        model, input_meta, output_meta = extract_model
        graph = gs_from_onnx(model)
        subgraph = extract_subgraph(graph, input_meta, output_meta)
        # Make sure original graph isn't modified.
        assert len(graph.nodes) == 2

        assert isinstance(subgraph, gs.Graph)
        assert len(subgraph.nodes) == 1

        assert len(subgraph.inputs) == 1
        assert subgraph.inputs[0].name == "X"

        assert len(subgraph.outputs) == 1
        assert subgraph.outputs[0].name == "identity_out_0"

    def test_extract_passes_no_input_shape(self, extract_model):
        model, input_meta, output_meta = extract_model
        input_meta["X"].shape = None
        model = extract_subgraph(model, input_meta, output_meta)
        self.check_model(model)

    def test_extract_passes_no_input_dtype(self, extract_model):
        model, input_meta, output_meta = extract_model
        input_meta["X"].dtype = None
        model = extract_subgraph(model, input_meta, output_meta)
        self.check_model(model)

    def test_extract_passes_no_output_shape(self, extract_model):
        model, input_meta, output_meta = extract_model
        output_meta["identity_out_0"].shape = None
        model = extract_subgraph(model, input_meta, output_meta)
        self.check_model(model)