From 89b22b64e522e359af3441fee845d6fef498e31e Mon Sep 17 00:00:00 2001 From: Jocelyn Date: Tue, 16 Mar 2021 15:31:10 -0400 Subject: [PATCH 1/3] add support for optional args for frontends tvmc --- python/tvm/driver/tvmc/frontends.py | 30 ++++++++++++---------- tests/python/driver/tvmc/test_frontends.py | 26 ++++++++++++++----- 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/python/tvm/driver/tvmc/frontends.py b/python/tvm/driver/tvmc/frontends.py index 16e6c8eb966e..a358a747d0a6 100644 --- a/python/tvm/driver/tvmc/frontends.py +++ b/python/tvm/driver/tvmc/frontends.py @@ -54,7 +54,7 @@ def suffixes(): """File suffixes (extensions) used by this frontend""" @abstractmethod - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): """Load a model from a given path. Parameters @@ -101,7 +101,7 @@ def name(): def suffixes(): return ["h5"] - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): # pylint: disable=C0103 tf, keras = import_keras() @@ -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") + kwargs["layout"] = layout + return relay.frontend.from_keras(model, input_shapes, **kwargs) def is_sequential_p(self, model): _, keras = import_keras() @@ -158,14 +160,14 @@ def name(): def suffixes(): return ["onnx"] - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): # pylint: disable=C0415 import onnx # pylint: disable=E1101 model = onnx.load(path) - return relay.frontend.from_onnx(model, shape=shape_dict) + return relay.frontend.from_onnx(model, shape=shape_dict, **kwargs) class TensorflowFrontend(Frontend): @@ -179,7 +181,7 @@ def name(): def suffixes(): return ["pb"] - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): # pylint: disable=C0415 import tensorflow as tf import tvm.relay.testing.tf as tf_testing @@ -192,7 +194,9 @@ def load(self, path, shape_dict=None): graph_def = tf_testing.ProcessGraphDefParam(graph_def) logger.debug("parse TensorFlow model and convert into Relay computation graph") - return relay.frontend.from_tensorflow(graph_def, shape=shape_dict) + return relay.frontend.from_tensorflow( + graph_def, shape=shape_dict, **kwargs + ) # doublecheck: nothing new needed class TFLiteFrontend(Frontend): @@ -206,7 +210,7 @@ def name(): def suffixes(): return ["tflite"] - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): # pylint: disable=C0415 import tflite.Model as model @@ -229,7 +233,7 @@ def load(self, path, shape_dict=None): raise TVMCException("input file not tflite version 3") logger.debug("parse TFLite model and convert into Relay computation graph") - mod, params = relay.frontend.from_tflite(tflite_model, shape_dict=shape_dict) + mod, params = relay.frontend.from_tflite(tflite_model, shape_dict=shape_dict, **kwargs) return mod, params @@ -245,7 +249,7 @@ def suffixes(): # Torch Script is a zip file, but can be named pth return ["pth", "zip"] - def load(self, path, shape_dict=None): + def load(self, path, shape_dict=None, **kwargs): # pylint: disable=C0415 import torch @@ -259,7 +263,7 @@ def load(self, path, shape_dict=None): input_shapes = list(shape_dict.items()) logger.debug("parse Torch model and convert into Relay computation graph") - return relay.frontend.from_pytorch(traced_model, input_shapes) + return relay.frontend.from_pytorch(traced_model, input_shapes, **kwargs) # whattest? ALL_FRONTENDS = [ @@ -339,7 +343,7 @@ def guess_frontend(path): raise TVMCException("failed to infer the model format. Please specify --model-format") -def load_model(path, model_format=None, shape_dict=None): +def load_model(path, model_format=None, shape_dict=None, **kwargs): """Load a model from a supported framework and convert it into an equivalent relay representation. @@ -367,6 +371,6 @@ def load_model(path, model_format=None, shape_dict=None): else: frontend = guess_frontend(path) - mod, params = frontend.load(path, shape_dict) + mod, params = frontend.load(path, shape_dict, **kwargs) return mod, params diff --git a/tests/python/driver/tvmc/test_frontends.py b/tests/python/driver/tvmc/test_frontends.py index b41f4c4dff2d..0fc15e0b6d6c 100644 --- a/tests/python/driver/tvmc/test_frontends.py +++ b/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): + mod, params = tvmc.frontends.load_model(model, **kwargs) assert type(mod) is IRModule assert type(params) is dict - ## check whether one known value is part of the params dict + return mod, params + + +def test_load_model__onnx(onnx_resnet50): + # some CI environments wont offer onnx, so skip in case it is not present + pytest.importorskip("onnx") + mod, params = verify_load_model__onnx(onnx_resnet50) + # check whether one known value is part of the params dict assert "resnetv24_batchnorm0_gamma" in params.keys() + mod, params = verify_load_model__onnx(onnx_resnet50, freeze_params=True) + # check that the parameter dict is empty, implying that they have been folded into constants + assert params == {} def test_load_model__pb(pb_mobilenet_v1_1_quant): From 9ef324f019fd66410880b2affc07aee7428f833c Mon Sep 17 00:00:00 2001 From: Jocelyn Date: Tue, 16 Mar 2021 15:36:05 -0400 Subject: [PATCH 2/3] remove unnecessary comments --- python/tvm/driver/tvmc/frontends.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/driver/tvmc/frontends.py b/python/tvm/driver/tvmc/frontends.py index a358a747d0a6..d474d9c3d73b 100644 --- a/python/tvm/driver/tvmc/frontends.py +++ b/python/tvm/driver/tvmc/frontends.py @@ -196,7 +196,7 @@ def load(self, path, shape_dict=None, **kwargs): logger.debug("parse TensorFlow model and convert into Relay computation graph") return relay.frontend.from_tensorflow( graph_def, shape=shape_dict, **kwargs - ) # doublecheck: nothing new needed + ) class TFLiteFrontend(Frontend): @@ -263,7 +263,7 @@ def load(self, path, shape_dict=None, **kwargs): input_shapes = list(shape_dict.items()) logger.debug("parse Torch model and convert into Relay computation graph") - return relay.frontend.from_pytorch(traced_model, input_shapes, **kwargs) # whattest? + return relay.frontend.from_pytorch(traced_model, input_shapes, **kwargs) ALL_FRONTENDS = [ From 02b33eec53fb4406348c4810bcee4832a96de3ba Mon Sep 17 00:00:00 2001 From: Jocelyn Date: Wed, 17 Mar 2021 15:26:58 -0400 Subject: [PATCH 3/3] Add changes suggested by Matt W. via PR --- python/tvm/driver/tvmc/frontends.py | 7 ++----- tests/python/driver/tvmc/test_frontends.py | 10 +++------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/python/tvm/driver/tvmc/frontends.py b/python/tvm/driver/tvmc/frontends.py index d474d9c3d73b..0488223c782f 100644 --- a/python/tvm/driver/tvmc/frontends.py +++ b/python/tvm/driver/tvmc/frontends.py @@ -130,8 +130,7 @@ def load(self, path, shape_dict=None, **kwargs): 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) - layout = kwargs.get("layout", "NHWC") - kwargs["layout"] = layout + kwargs.setdefault("layout", "NHWC") return relay.frontend.from_keras(model, input_shapes, **kwargs) def is_sequential_p(self, model): @@ -194,9 +193,7 @@ def load(self, path, shape_dict=None, **kwargs): graph_def = tf_testing.ProcessGraphDefParam(graph_def) logger.debug("parse TensorFlow model and convert into Relay computation graph") - return relay.frontend.from_tensorflow( - graph_def, shape=shape_dict, **kwargs - ) + return relay.frontend.from_tensorflow(graph_def, shape=shape_dict, **kwargs) class TFLiteFrontend(Frontend): diff --git a/tests/python/driver/tvmc/test_frontends.py b/tests/python/driver/tvmc/test_frontends.py index 0fc15e0b6d6c..5a63c5c47933 100644 --- a/tests/python/driver/tvmc/test_frontends.py +++ b/tests/python/driver/tvmc/test_frontends.py @@ -115,22 +115,18 @@ def test_load_model__tflite(tflite_mobilenet_v1_1_quant): assert "_param_1" in params.keys() -def verify_load_model__keras(model, **kwargs): +@pytest.mark.parametrize("load_model_kwargs", [{}, {"layout": "NCHW"}]) +def test_load_model__keras(keras_resnet50, load_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(model) + mod, params = tvmc.frontends.load_model(keras_resnet50, **load_model_kwargs) 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__keras(keras_resnet50): - verify_load_model__keras(keras_resnet50) - verify_load_model__keras(keras_resnet50, layout="NCHW") - - def verify_load_model__onnx(model, **kwargs): mod, params = tvmc.frontends.load_model(model, **kwargs) assert type(mod) is IRModule