mbaret commented on a change in pull request #5634: URL: https://github.com/apache/incubator-tvm/pull/5634#discussion_r430255016
########## File path: tests/python/contrib/test_coreml_codegen.py ########## @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 pytest +import numpy as np + +import tvm +from tvm import relay +from tvm.relay import transform +from tvm.contrib.target import coreml as _coreml + +def check_result(mod, map_inputs, out_shape, result, tol=1e-3, target="llvm", + ctx=tvm.cpu(), params=None): + def check_graph_runtime_result(): + with relay.build_config(opt_level=3): + json, lib, param = relay.build(mod, target=target, params=params) + rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx) + + for name, data in map_inputs.items(): + rt_mod.set_input(name, data) + rt_mod.set_input(**param) + rt_mod.run() + out = tvm.nd.empty(out_shape, ctx=ctx) + out = rt_mod.get_output(0, out) + + tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol) + + check_graph_runtime_result() + [email protected]('skip because coremltools is not available in CI') Review comment: Does this mean there will be no testing in CI for this change? I think we should find a way to test as much as possible or enable CI to perform some proper tests, otherwise we're not going to have confidence we don't break this integration. For instance, is there a way we could just test the compilation pathway? ########## File path: python/tvm/contrib/target/coreml.py ########## @@ -0,0 +1,252 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel +"""Utility to compile CoreML models""" + +import os +import shutil + +import tvm._ffi +from ...relay.expr_functor import ExprVisitor +from ...relay.expr import Constant +from ...relay import op as _op +from .. import xcode, coreml_runtime + +def _convert_add(builder, name, inputs, outputs, args, attrs): + builder.add_elementwise( + name=name, + input_names=inputs, + output_name=outputs[0], + mode='ADD' + ) + +def _convert_multiply(builder, name, inputs, outputs, args, attrs): + builder.add_elementwise( + name=name, + input_names=inputs, + output_name=outputs[0], + mode='MULTIPLY' + ) + +def _convert_clip(builder, name, inputs, outputs, args, attrs): + builder.add_clip( + name=name, + input_name=inputs[0], + output_name=outputs[0], + min_value=attrs.a_min, + max_value=attrs.a_max + ) + +def _convert_batch_flatten(builder, name, inputs, outputs, args, attrs): + builder.add_flatten_to_2d( + name=name, + input_name=inputs[0], + output_name=outputs[0] + ) + +def _convert_softmax(builder, name, inputs, outputs, args, attrs): + builder.add_softmax_nd( + name=name, + input_name=inputs[0], + output_name=outputs[0], + axis=int(attrs['axis']) + ) + +def _convert_conv2d(builder, name, inputs, outputs, args, attrs): + weight = args[1].data.asnumpy() + if attrs['kernel_layout'] == 'OIHW': + # convert to 'HWIO' + weight = weight.transpose([2, 3, 1, 0]) + kh, kw, kc, oc = weight.shape + + builder.add_convolution( + name=name, + kernel_channels=kc, + output_channels=oc, + height=kh, + width=kw, + stride_height=int(attrs['strides'][0]), + stride_width=int(attrs['strides'][0]), + border_mode="valid", + groups=int(attrs['groups']), + W=weight, + b=None, + has_bias=False, + input_name=inputs[0], + output_name=outputs[0], + dilation_factors=[int(v) for v in attrs['dilation']], + padding_top=int(attrs['padding'][0]), + padding_bottom=int(attrs['padding'][2]), + padding_left=int(attrs['padding'][1]), + padding_right=int(attrs['padding'][3]) + ) + +def _convert_global_avg_pool2d(builder, name, inputs, outputs, args, attrs): + builder.add_pooling( + name=name, + height=1, + width=1, + stride_height=1, + stride_width=1, + layer_type='AVERAGE', + padding_type='VALID', + input_name=inputs[0], + output_name=outputs[0], + is_global=True + ) + +_convert_map = { + 'add' : _convert_add, + 'multiply' : _convert_multiply, + 'clip' : _convert_clip, + 'nn.batch_flatten' : _convert_batch_flatten, + 'nn.softmax' : _convert_softmax, + 'nn.conv2d' : _convert_conv2d, + 'nn.global_avg_pool2d' : _convert_global_avg_pool2d, +} + +class CodegenCoreML(ExprVisitor): + """ + A visitor to traverse subgraphs and build Core ML models. + """ + def __init__(self, model_name, function): + import coremltools + from coremltools.models.neural_network import NeuralNetworkBuilder + + ExprVisitor.__init__(self) + self.model_name = model_name + self.function = function + self.out_map = {} + self.model_inputs_ = [] + self.buf_idx_ = 0 + + # Update inputs and outputs after we visit all the nodes. + # Set dummy values for now. + # TODO: support multiple outputs + inputs = [('', coremltools.models.datatypes.Array(1,)) for _ in self.function.params] + outputs = [('', coremltools.models.datatypes.Array(1,))] + self.builder = NeuralNetworkBuilder(inputs, outputs, + disable_rank5_shape_mapping=True) + + def visit_constant(self, const): + output = "buf_" + str(self.buf_idx_) + self.builder.add_load_constant_nd( + name=output, + output_name=output, + constant_value=const.data.asnumpy(), + shape=const.data.shape + ) + self.buf_idx_ = self.buf_idx_ + 1 + self.out_map[const] = [output] + + def visit_var(self, var): + name = var.name_hint + shape = [int(n) for n in var.type_annotation.shape] + dtype = var.type_annotation.dtype + self.model_inputs_.append((name, shape, dtype)) + self.out_map[var] = [name] + + def visit_call(self, call): + inputs = [] + for arg in call.args: + super().visit(arg) + for out in self.out_map[arg]: + inputs.append(out) + outputs = ["buf_" + str(self.buf_idx_)] + op_name = call.op.name + layer_name = op_name + "_" + str(self.buf_idx_) + + assert op_name in _convert_map, "{} is not supported".format(op_name) + _convert_map[op_name](self.builder, layer_name, inputs, outputs, + call.args, call.attrs) + + self.buf_idx_ = self.buf_idx_ + 1 + self.out_map[call] = outputs + + def compile(self, out_dir): + """ + Build a Core ML model and compile it with Xcode toolchain. + """ + import coremltools + from coremltools.proto.Model_pb2 import ArrayFeatureType + + FEATURE_TYPE_MAP = { + "float32": ArrayFeatureType.FLOAT32, + "float64": ArrayFeatureType.DOUBLE, + "int32": ArrayFeatureType.INT32, + } + + input_names, input_dims, input_dtypes = zip(*self.model_inputs_) + self.builder.set_input(input_names, input_dims) + for i, dtype in enumerate(input_dtypes): + assert dtype in FEATURE_TYPE_MAP + input_desc = self.builder.spec.description.input + input_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype] + + output_dim = [int(n) for n in self.function.ret_type.shape] + self.builder.set_output(self.out_map[self.function.body], [output_dim]) + for i, dtype in enumerate([self.function.ret_type.dtype]): + assert dtype in FEATURE_TYPE_MAP + output_desc = self.builder.spec.description.output + output_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype] + + model = coremltools.models.MLModel(self.builder.spec) + xcode.compile_coreml(model, self.model_name, out_dir) + + +@tvm._ffi.register_func("relay.ext.coremlcompiler") +def coreml_compiler(ref): + """ + Create a CoreML runtime from a Relay module. + """ + model_dir = os.getcwd() + if isinstance(ref, tvm.ir.module.IRModule): + for var, func in ref.functions.items(): + name = var.name_hint + builder = CodegenCoreML(name, func) + builder.visit(func.body) + mlmodelc_path = "{}/{}.mlmodelc".format(model_dir, name) + if os.path.exists(mlmodelc_path): + shutil.rmtree(mlmodelc_path) + builder.compile(model_dir) + + ctx = tvm.cpu(0) + return coreml_runtime.create(model_dir, ctx).module + + +def _register_coreml_op(op_name): Review comment: This should go in python/tvm/relay/op/contrib/coreml.py (you can take a look at dnnl.py in the same directory). ########## File path: tests/python/contrib/test_coreml_codegen.py ########## @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 pytest +import numpy as np + +import tvm +from tvm import relay +from tvm.relay import transform +from tvm.contrib.target import coreml as _coreml + +def check_result(mod, map_inputs, out_shape, result, tol=1e-3, target="llvm", + ctx=tvm.cpu(), params=None): + def check_graph_runtime_result(): + with relay.build_config(opt_level=3): + json, lib, param = relay.build(mod, target=target, params=params) + rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx) + + for name, data in map_inputs.items(): + rt_mod.set_input(name, data) + rt_mod.set_input(**param) + rt_mod.run() + out = tvm.nd.empty(out_shape, ctx=ctx) + out = rt_mod.get_output(0, out) + + tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol) + + check_graph_runtime_result() + [email protected]('skip because coremltools is not available in CI') +def test_coreml_codegen(): + shape = (1,) + x = relay.var('x', shape=shape) + y = relay.var('y', shape=shape) + _x = relay.annotation.compiler_begin(x, "coremlcompiler") + _y = relay.annotation.compiler_begin(y, "coremlcompiler") + z = _x + _x + p = _y * _y + p = relay.annotation.compiler_end(p, "coremlcompiler") + z = relay.annotation.compiler_end(z, "coremlcompiler") + f = relay.Function([x, y], p - z) + x_data = np.random.rand(*shape).astype('float32') + y_data = np.random.rand(*shape).astype('float32') + mod = tvm.IRModule() + mod["main"] = f + mod = transform.PartitionGraph()(mod) Review comment: I think it would be best to test the entire BYOC flow here instead of using manual annotation. ---------------------------------------------------------------- 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: [email protected]
