[GitHub] [tvm] mdw-octoml commented on a change in pull request #7674: [TVMC] Allow optional arguments to be passed to importers

[GitBox <gi...@apache.org>]

mdw-octoml commented on a change in pull request #7674:
URL: https://github.com/apache/tvm/pull/7674#discussion_r595513989



##########
File path: tests/python/driver/tvmc/test_frontends.py
##########
@@ -115,26 +115,38 @@ def test_load_model__tflite(tflite_mobilenet_v1_1_quant):
     assert "_param_1" in params.keys()
 
 
-def test_load_model__keras(keras_resnet50):
+def verify_load_model__keras(model, **kwargs):
     # some CI environments wont offer TensorFlow/Keras, so skip in case it is not present
     pytest.importorskip("tensorflow")
 
-    mod, params = tvmc.frontends.load_model(keras_resnet50)
+    mod, params = tvmc.frontends.load_model(model)

Review comment:
       I think you need to pass kwargs here, right?

##########
File path: python/tvm/driver/tvmc/frontends.py
##########
@@ -130,7 +130,9 @@ def load(self, path, shape_dict=None):
         input_shapes = {name: x.shape for (name, x) in zip(model.input_names, inputs)}
         if shape_dict is not None:
             input_shapes.update(shape_dict)
-        return relay.frontend.from_keras(model, input_shapes, layout="NHWC")
+        layout = kwargs.get("layout", "NHWC")

Review comment:
       To make this more clear I would do `kwargs.setdefault("layout", "NHWC")` instead -- this does the same thing (I think!)

##########
File path: tests/python/driver/tvmc/test_frontends.py
##########
@@ -115,26 +115,38 @@ def test_load_model__tflite(tflite_mobilenet_v1_1_quant):
     assert "_param_1" in params.keys()
 
 
-def test_load_model__keras(keras_resnet50):
+def verify_load_model__keras(model, **kwargs):
     # some CI environments wont offer TensorFlow/Keras, so skip in case it is not present
     pytest.importorskip("tensorflow")
 
-    mod, params = tvmc.frontends.load_model(keras_resnet50)
+    mod, params = tvmc.frontends.load_model(model)
     assert type(mod) is IRModule
     assert type(params) is dict
     ## check whether one known value is part of the params dict
     assert "_param_1" in params.keys()
 
 
-def test_load_model__onnx(onnx_resnet50):
-    # some CI environments wont offer onnx, so skip in case it is not present
-    pytest.importorskip("onnx")
+def test_load_model__keras(keras_resnet50):

Review comment:
       A more Pythonic way to do this is to have a single test case (`test_load_model__keras`) and use `@pytest.mark.parametrize(...)` to specify the different test cases. In this case you will want to pass the full parameters to load_model as a single `kwargs` dict, so it would be something like:
   
   ```
   @pytest.mark.parametrize("load_model_kwargs", [ {"path": keras_resnet50}, {"path": keras_resnet50, "layout": "NCHW"}])
   def test_load_model__keras(load_model_kwargs):
     pytest.importorskip("tensorflow")
     mod, params = tvmc.frontends.load_model(load_model_kwargs)
     # ... rest of the test here...
   ```
   
   This also means you can fold `verify_load_model__keras` back into the test rather than having a helper function.

##########
File path: tests/python/driver/tvmc/test_frontends.py
##########
@@ -115,26 +115,38 @@ def test_load_model__tflite(tflite_mobilenet_v1_1_quant):
     assert "_param_1" in params.keys()
 
 
-def test_load_model__keras(keras_resnet50):
+def verify_load_model__keras(model, **kwargs):
     # some CI environments wont offer TensorFlow/Keras, so skip in case it is not present
     pytest.importorskip("tensorflow")
 
-    mod, params = tvmc.frontends.load_model(keras_resnet50)
+    mod, params = tvmc.frontends.load_model(model)
     assert type(mod) is IRModule
     assert type(params) is dict
     ## check whether one known value is part of the params dict
     assert "_param_1" in params.keys()
 
 
-def test_load_model__onnx(onnx_resnet50):
-    # some CI environments wont offer onnx, so skip in case it is not present
-    pytest.importorskip("onnx")
+def test_load_model__keras(keras_resnet50):
+    verify_load_model__keras(keras_resnet50)
+    verify_load_model__keras(keras_resnet50, layout="NCHW")
+
 
-    mod, params = tvmc.frontends.load_model(onnx_resnet50)
+def verify_load_model__onnx(model, **kwargs):

Review comment:
       As in the comment above, I would probably suggest using `@pytest.mark.parametrize` here as well and not having the separate `verify_load_model__onnx` helper function.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org