Merge pull request #339 from aseemw/dev/add_const

Added squeeze, unsqueeze, const conversions

onnx_coreml/_graph.py

@@ -139,6 +139,8 @@ def __init__(self,
  # data blob name to the op_type that generates it
  self.blob_from_op_type = {} # type: Dict[Text, Text]
 
+ self.constant_layers_added = {} # type: Dict[Text, bool]
+
  for node_ in nodes:
  for input_ in node_.inputs:
  if input_ in self.blob_to_op_type:

onnx_coreml/_operators.py

@@ -11,8 +11,6 @@
 from coremltools.proto import NeuralNetwork_pb2 #type: ignore
 from ._error_utils import ErrorHandling
 
-_SEQUENCE_LAYERS_REGISTRY = set(["LSTM"])
-
 def _compare(a, b, encoding="utf8"): #type: (Text, Text, Text) -> bool
  if isinstance(a, bytes):
  a = a.decode(encoding)
@@ -348,9 +346,6 @@ def _convert_add(builder, node, graph, err): # type: (NeuralNetworkBuilder, Nod
  shape_bias=[second_input.shape[0]])
  return
 
- if 'broadcast' in node.attrs:
- if node.attrs['broadcast'] == 1:
- return err.unsupported_op_configuration(builder, node, graph, "Broadcast Add is not supported now")
  builder.add_elementwise(
  name=node.name,
  input_names=node.inputs,
@@ -359,10 +354,6 @@ def _convert_add(builder, node, graph, err): # type: (NeuralNetworkBuilder, Nod
  )
 
 def _convert_mul(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
- if 'broadcast' in node.attrs:
- if node.attrs['broadcast'] == 1:
- return err.unsupported_op_configuration(builder, node, graph, "Broadcast Multiply is not supported now")
-
  builder.add_elementwise(
  name=node.name,
  input_names=node.inputs,
@@ -371,10 +362,6 @@ def _convert_mul(builder, node, graph, err): # type: (NeuralNetworkBuilder, Nod
  )
 
 def _convert_div(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
- if 'broadcast' in node.attrs:
- if node.attrs['broadcast'] == 1:
- return err.unsupported_op_configuration(builder, node, graph, "Broadcast Div is not supported now")
-
  builder.add_unary(name=node.name + '_inverse', #type: ignore
  input_name=node.inputs[1],
  output_name=node.inputs[1] + '_inverse',
@@ -985,6 +972,34 @@ def _convert_custom(builder, node, graph, err): # type: (NeuralNetworkBuilder, N
 
  err.custom_layer_nodes.append(node)
 
+def _convert_identity(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
+ builder.add_activation(
+ name=node.name,
+ non_linearity = 'LINEAR',
+ input_name=node.inputs[0],
+ output_name=node.outputs[0],
+ params=[1.0, 0.0]
+ )
+
+def _convert_const(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
+
+ for name, value in node.input_tensors.items(): 
+ if name not in graph.constant_layers_added:
+ shape = value.shape
+ coreml_shape = [1,1,1]
+ if len(shape) == 3:
+ coreml_shape = list(shape)
+ elif len(shape) == 1:
+ coreml_shape = [shape[0],1,1]
+ elif len(shape) == 2:
+ coreml_shape = [1, shape[0], shape[1]]
+ else:
+ return err.unsupported_op_configuration(builder, node, graph, "unable to translate constant array shape to CoreML shape")
+ builder.add_load_constant(name=name,
+ output_name=name,
+ constant_value=value.flatten(),
+ shape=coreml_shape)
+ graph.constant_layers_added[name] = True
 
 
 _ONNX_NODE_REGISTRY = {
@@ -1050,8 +1065,13 @@ def _convert_custom(builder, node, graph, err): # type: (NeuralNetworkBuilder, N
  "ArgMin": _convert_reduce,
  "Clip": _convert_clip,
  "MeanVarianceNormalization": _convert_mvn,
+ "Unsqueeze": _convert_identity,
+ "Squeeze": _convert_identity
 }
 
+_SEQUENCE_LAYERS_REGISTRY = set(["LSTM"])
+
+_CONST_INPUT_ALLOWED_LAYERS = set([ "Add", "Sum", "Mul", "Concat", "Max", "Min", "Div", "Reciprocal"])
 
 def _get_node_converter_fn(builder, node, err): # type: (NeuralNetworkBuilder, Node, ErrorHandling) -> Callable[[NeuralNetworkBuilder, Node, Graph, ErrorHandling], None]
  """
@@ -1063,6 +1083,12 @@ def _get_node_converter_fn(builder, node, err): # type: (NeuralNetworkBuilder,
  else:
  return err.unsupported_op(node)
 
+def _add_const_inputs_if_required(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
+ if node.op_type in _CONST_INPUT_ALLOWED_LAYERS:
+ if len(node.input_tensors) > 0:
+ _convert_const(builder, node, graph, err)
+
+
 def _convert_node(builder, node, graph, err): # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
  converter_fn = _get_node_converter_fn(builder, node, err)
  return converter_fn(builder, node, graph, err)

onnx_coreml/converter.py

@@ -16,7 +16,7 @@
 
 from typing import Tuple
 
-from ._operators import _convert_node, _SEQUENCE_LAYERS_REGISTRY, _ONNX_NODE_REGISTRY
+from ._operators import _convert_node, _SEQUENCE_LAYERS_REGISTRY, _ONNX_NODE_REGISTRY, _add_const_inputs_if_required
 from ._graph import Graph, EdgeInfo, Transformer
 from ._transformers import ConvAddFuser, DropoutRemover, \
  ReshapeInitTensorFuser, BNBroadcastedMulFuser, BNBroadcastedAddFuser, \
@@ -410,6 +410,7 @@ def convert(model, # type: Union[onnx.ModelProto, Text]
 
  for i, node in enumerate(graph.nodes):
  print("%d/%d: Converting Node Type %s" %(i+1, len(graph.nodes), node.op_type))
+ _add_const_inputs_if_required(builder, node, graph, err)
  _convert_node(builder, node, graph, err)
 
  if add_deprocess:
@@ -460,7 +461,12 @@ def convert(model, # type: Union[onnx.ModelProto, Text]
  if outputs.name == output_:
  builder.spec.description.output[i].shortDescription = 'This output is a sequence'
 
- mlmodel = MLModel(builder.spec)
+ print("Translation to CoreML spec completed. Now compiling the CoreML model.")
+ try:
+ mlmodel = MLModel(builder.spec)
+ except:
+ raise ValueError('Compilation failed. Translation to CoreML spec was incorrect.')
+
 
  # print information about all ops for which custom layers have been added
  if len(err.custom_layer_nodes) > 0:

tests/model_test.py

@@ -41,7 +41,7 @@ def _test_torch_model_single_io(torch_model, torch_input_shape, coreml_input_sha
 
  # delete onnx model
  if os.path.exists(model_dir):
- shutil.rmtree(model_dir)
+  shutil.rmtree(model_dir)
 
 class OnnxModelTest(unittest.TestCase):
 
@@ -77,6 +77,35 @@ def forward(self, x):
  torch_model.train(False)
  _test_torch_model_single_io(torch_model, (1, 3, 100, 100), (3, 100, 100)) # type: ignore
 
+ def test_const_initializer1(self): # typr: () -> None
+ class Net(nn.Module):
+ def __init__(self):
+ super(Net, self).__init__()
+ self.ones = torch.nn.Parameter(torch.ones(1,))
+
+ def forward(self, x):
+ y = x + self.ones
+ return y
+
+ torch_model = Net() # type: ignore
+ torch_model.train(False)
+ _test_torch_model_single_io(torch_model, (1, 3), (3,)) # type: ignore
+
+
+ def test_const_initializer2(self): # typr: () -> None
+ class Net(nn.Module):
+ def __init__(self):
+ super(Net, self).__init__()
+
+ def forward(self, x):
+ y = x + torch.nn.Parameter(torch.ones(2, 3))
+ return y
+
+ torch_model = Net() # type: ignore
+ torch_model.train(False)
+ _test_torch_model_single_io(torch_model, (1, 2, 3), (1, 2, 3)) # type: ignore
+
+
 
 if __name__ == '__main__':
  unittest.main()

tests/onnx_backend_node_test.py

@@ -127,11 +127,9 @@ def run_node(cls,
 backend_test.exclude('test_log_softmax_lastdim_cpu')
 backend_test.exclude('test_softmax_functional_dim3_cpu')
 backend_test.exclude('test_softmax_lastdim_cpu')
-backend_test.exclude('test_squeeze_cpu')
 backend_test.exclude('test_sub_bcast_cpu')
 backend_test.exclude('test_sub_cpu')
 backend_test.exclude('test_sub_example_cpu')
-backend_test.exclude('test_unsqueeze_cpu')
 backend_test.exclude('test_slice_end_out_of_bounds_cpu')
 backend_test.exclude('test_slice_neg_cpu')
 backend_test.exclude('test_GLU_cpu')