Skip to content

Commit

Permalink
[Coreml] Fix Coreml Input Shape Handling (apache#8562)
Browse files Browse the repository at this point in the history
* convert ot python list like expected

* test example

* jostle ci

Co-authored-by: Andrew Zhao Luo <andrewzhaoluo@system76-pc.localdomain>
  • Loading branch information
2 people authored and ylc committed Sep 29, 2021
1 parent 157fe48 commit 3c7e561
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 1 deletion.
2 changes: 1 addition & 1 deletion python/tvm/relay/frontend/coreml.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,7 +562,7 @@ def from_coreml(model, shape=None):

etab = ExprTable()
for i in spec.description.input:
input_shape = shape[i.name] if shape is not None and i.name in shape else None
input_shape = list(shape[i.name]) if shape is not None and i.name in shape else None
etab.set_expr(i.name, _expr.var(i.name, shape=input_shape))

for pp in cc.preprocessing:
Expand Down
44 changes: 44 additions & 0 deletions tests/python/frontend/coreml/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@
import coremltools as cm
import model_zoo
import tvm.testing
import tempfile
from os import path
import tvm
import tvm.relay as relay
import tensorflow.keras as keras


def get_tvm_output(
Expand Down Expand Up @@ -783,6 +788,44 @@ def test_forward_convolution():
verify_convolution((1, 3, 224, 224), filter=(32, 3, 3, 3), padding="SAME")


def test_can_build_keras_to_coreml_to_relay():
"""Test multiple conversion paths and importing from
a saved file."""
model = keras.models.Sequential()
model.add(
keras.layers.Conv2D(
filters=6,
kernel_size=(1, 1),
activation="relu",
padding="same",
input_shape=(3, 3, 1),
data_format="channels_first",
)
)

with tempfile.TemporaryDirectory() as tmpdir:
kmodel_fn = path.join(tmpdir, "c1mdl.h5")
model.save(kmodel_fn)

mdl = cm.convert(kmodel_fn)
model_file = path.join(tmpdir, "c1.mlmodel")
mdl.save(model_file)

mdl = cm.models.MLModel(model_file)
desc = mdl.get_spec().description
iname = desc.input[0].name
ishape = desc.input[0].type.multiArrayType.shape
shape_dict = {}
for i in mdl.get_spec().description.input:
iname = i.name
ishape = i.type.multiArrayType.shape
shape_dict[iname] = ishape
mod, params = relay.frontend.from_coreml(mdl, shape_dict)

with tvm.transform.PassContext(opt_level=3):
relay.build(mod, "llvm", params=params)


if __name__ == "__main__":
test_forward_AddLayerParams()
test_forward_ConcatLayerParams()
Expand All @@ -801,3 +844,4 @@ def test_forward_convolution():
test_resnet50_checkonly()
test_forward_image_scaler()
test_forward_convolution()
test_can_build_keras_to_coreml_to_relay()

0 comments on commit 3c7e561

Please sign in to comment.