ONNX test code cleanup

[vandanavk]

Description

This PR refactors ONNX import/export test code

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

Current directory structure:

Redundant code:

1. test_case list - The same test case list (almost) is present in both export and import
2. GluonBackendRep is available in import alone. This may become redundant when Gluon→ONNX export is added
3. mxnet_backend_test, onnx_backend_test, gluon_backend_test all have the same set of loops
4. mxnet_backend.py, backend.py - almost similar code
5. mxnet_export_test, onnx_import_test - test_models is redundant, there is scope for combining operator tests.

Changes that are part of cleanup:

• Create tests/python-pytest/onnx/test_cases.py

IMPLEMENTED_OPERATORS_TEST = {'both' : [...], 'import': [...], 'export' : [...]}
BASIC_MODEL_TESTS = {'both' : [...], 'import': [...], 'export' : [...]}
STANDARD_MODEL = {'both' : [...], 'import': [...], 'export' : [...]}

• Create onnx/mxnet_backend_test.py and onnx/gluon_backend_test.py for import/export with MXNet/Gluon backend respectively. These can't be combined further because the ONNX tests are added in globals() dictionary with test name as key. Combining these 2 will override the test list in globals().

• Reference this newly created test_cases.py in mxnet_backend_test.py, onnx_backend_test.py, gluon_backend_test.py. Create a common onnx/backend_test.py for preparing the test list.

def prepare_tests(backend, operation):
    """
    Prepare the test list
    :param backend: mxnet/gluon backend
    :param operation: str. export or import
    :return: backend test list
    """
    ...
    ...

    return BACKEND_TESTS

Example call:

BACKEND_TESTS = backend_test.prepare_tests(mxnet_backend, 'import')

• Instead of having separate MXNetBackend for MXNet import/export, GluonBackend for Gluon import, merge into one by either passing special params to prepare() or setting class variables

class MXNetBackend(Backend):
    """MXNet/Gluon backend for ONNX"""

    backend = 'mxnet'
    operation = 'import'

    @classmethod
    def set_params(cls, backend, operation):
        cls.backend = backend
        cls.operation = operation

@classmethod
    def prepare(cls, model, device='CPU', **kwargs):
        """For running end to end model(used for onnx test backend)

        Parameters
        ----------
        model  : onnx ModelProto object
            loaded onnx graph
        device : 'CPU'
            specifying device to run test on
        kwargs :
            other arguments

        Returns
        -------
        MXNetBackendRep : object
            Returns object of MXNetBackendRep class which will be in turn
            used to run inference on the input model and return the result for comparison.
        """
        backend = kwargs.get('backend', cls.backend)
        operation = kwargs.get('operation', cls.operation)
        ...

Option 1:

mxnet_backend.MXNetBackend.set_params('mxnet', 'export')

Option 2:

bkd_rep = backend.prepare(model, backend='mxnet', operation='export')

• Move GluonBackendRep to onnx/backend_rep.py and access from Gluon backend.
• Create a separate file onnx/test_models.py.

class TestModel(unittest.TestCase):
    """ Tests for models.
    Tests are dynamically added.
    Therefore edit test_models to add more tests.
    """
    def test_import_export(self):
        for test in test_cases:
            model_name, input_shape, output_shape = test
            with self.subTest(model_name):
                ...
                ...

To add tests for models, add a line in test_cases[]

test_cases = [
    (model name, input shape, output shape)
]

Example:

test_cases = [
    ("bvlc_googlenet", (1, 3, 224, 224), (1, 1000))
]

• Create onnx/test_node.py for custom tests for operators (which are specific to MXNet or for which there are no tests in ONNX)

class TestNode(unittest.TestCase):
    """ Custom tests for operators.
    Tests are dynamically added.
    Therefore edit test_cases to add more tests.
    """
    def test_import_export(self):
        for test in test_cases:
            test_name, mxnet_op, onnx_name, inputs, attrs = test
            with self.subTest(test_name):
                ...
                ...

To add tests for operators,

# test_case = ("test_case_name", mxnet op, "ONNX_op_name", [input_list], attribute map, MXNet specific operator = True/False)
test_cases = [
    ("test_equal", mx.sym.broadcast_equal, "Equal", [get_rnd((1, 3, 4, 5)), get_rnd((1, 5))], {}, False),
    ("test_softmax", mx.sym.SoftmaxOutput, "Softmax", [get_rnd((1000, 1000)), get_rnd(1000)],
     {'ignore_label': 0, 'use_ignore': False}, True)

• mxnet_export_test.py has all tests for MXNet to ONNX export

class TestExport(unittest.TestCase):
    """ Tests ONNX export."""
    def test_onnx_export_single_output(self):
        ...
    def test_onnx_export_multi_output(self):
        ...
    def test_onnx_export_list_shape(self):
        ...
    def test_onnx_export_extra_params(self):
        ...

Directory structure after cleanup:

Future work:

• mxnet_backend_test, onnx_backend_test, gluon_backend_test could be merged together by writing our own test cases (/~https://github.com/onnx/onnx-tensorflow/blob/master/test/frontend/test_node.py), instead of ONNX tests

Comments

• CI script has been updated

[vandanavk]

@mxnet-label-bot update [ONNX, pr-awaiting-review]

@Roshrini @anirudhacharya @zhreshold for review

[vandanavk]

Adding @safrooze to review onnx/mxnet_export_test.py (Ref: #13390).

Adding @nswamy to review onnx/test_models.py and onnx/test_node.py (Ref: #12852 (comment))

[sandeep-krishnamurthy]

Thanks.
Excellent PR description!!

[sandeep-krishnamurthy]

@nswamy @anirudhacharya @safrooze - Going ahead with merging this PR as it contains nice clean up in test cases and with approvals from other community members.