Adding ML Model passes

Author: bhushan23

onnx_coreml/converter.py

@@ -28,7 +28,11 @@
     ReshapeInitTensorFuser, BNBroadcastedMulFuser, BNBroadcastedAddFuser, \
     PixelShuffleFuser, OutputRenamer, AddModelInputsOutputs, \
     ConstantsToInitializers, ImageScalerRemover, ShapeOpRemover, ConstantRemover, \
-    ConstantFillToInitializers, ReshapeTransposeReshape_pattern1, CastOpRemover, DeadCodeElimination
+    ConstantFillToInitializers, ReshapeTransposeReshape_pattern1, CastOpRemover, \
+    DeadCodeElimination
+
+# ML model passes
+from coremltools.converters.nnssa.coreml.graph_pass.mlmodel_passes import remove_disconnected_constants
 
 from ._error_utils import ErrorHandling
 from .graph_viz import plot_graph # type: ignore
@@ -717,6 +721,11 @@ def _add_informative_description(feature, raise_error=True):
             if layer.WhichOneof('layer') == 'resizeBilinear' or layer.WhichOneof('layer') == 'cropResize':
                 raise TypeError('{} not supported with target iOS 11.2 please provide higher target iOS'.format(layer.WhichOneof('layer')))
 
+    # Optimize ML Model Spec
+    ml_model_passes = [remove_disconnected_constants]
+    for opt in ml_model_passes:
+        opt(builder.spec)
+
     print("Translation to CoreML spec completed. Now compiling the CoreML model.")
     try:
         if DEBUG:

tests/test_mlmodel_passes.py

@@ -0,0 +1,32 @@
+import numpy as np
+import unittest
+import coremltools.models.datatypes as datatypes
+from coremltools.models import neural_network as neural_network
+from coremltools.converters.nnssa.coreml.graph_pass.mlmodel_passes import remove_disconnected_constants
+
+
+class MLModelPassesTest(unittest.TestCase):
+
+    def test_load_constant_remove(self):
+        input_features = [('data', datatypes.Array(*(3, 4)))]
+        output_features = [('out', None)]
+        builder = neural_network.NeuralNetworkBuilder(input_features, output_features, disable_rank5_shape_mapping=True)
+        builder.add_activation('relu1', 'RELU', 'data', 'relu1')
+        builder.add_load_constant_nd('const1', 'c1', constant_value=np.ones((5,)), shape=(5,))
+        builder.add_activation('relu2', 'RELU', 'relu1', 'out')
+        builder.add_load_constant_nd('const2', 'c2', constant_value=np.ones((5,)), shape=(5,))
+        builder.add_load_constant_nd('const3', 'c3', constant_value=np.ones((5,)), shape=(5,))
+        spec = builder.spec
+        np.testing.assert_equal(5, len(spec.neuralNetwork.layers))
+        remove_disconnected_constants(spec)
+        np.testing.assert_equal(2, len(spec.neuralNetwork.layers))
+
+
+if __name__ == '__main__':
+    RUN_ALL_TESTS = True
+    if RUN_ALL_TESTS:
+        unittest.main()
+    else:
+        suite = unittest.TestSuite()
+        suite.addTest(MLModelPassesTest('test_load_constant_remove'))
+        unittest.TextTestRunner().run(suite)