Convert Keras Separable Convolution to Onnx (#201)

* Separable Convolution of keras #197

Author: jiafatom

onnxmltools/convert/keras/operator_converters/Conv.py

@@ -6,9 +6,11 @@
 import numpy
 import keras
 from distutils.version import StrictVersion
-from keras.layers import Conv1D, Conv2D, Conv3D, Conv2DTranspose, Conv3DTranspose
+from keras.layers import Conv1D, Conv2D, Conv3D, Conv2DTranspose, Conv3DTranspose, SeparableConv2D
 if StrictVersion(keras.__version__) >= StrictVersion('2.1.5'):
     from keras.layers import DepthwiseConv2D
+if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+    from keras.layers import SeparableConv1D
 from ....proto import onnx_proto
 from ...common._apply_operation import apply_identity, apply_transpose
 from ...common._registration import register_converter
@@ -32,10 +34,42 @@ def _calc_explicit_padding(input_size, output_shape, output_padding, kernel_shap
     return pads
 
 
+def process_separable_conv_2nd(scope, operator, container, convolution_input_names, n_dims,
+                               weight_perm_axes, parameters, auto_pad):
+    attrs = {'name': operator.full_name + '1'}
+
+    weight_tensor_name = scope.get_unique_variable_name('W')
+    weight_params = parameters[1].transpose(weight_perm_axes)
+    container.add_initializer(weight_tensor_name, onnx_proto.TensorProto.FLOAT,
+                              weight_params.shape, weight_params.flatten())
+    convolution_input_names.append(weight_tensor_name)
+
+    if len(parameters) == 3:
+        bias_tensor_name = scope.get_unique_variable_name('B')
+        container.add_initializer(bias_tensor_name, onnx_proto.TensorProto.FLOAT,
+                                  parameters[2].shape, parameters[2].flatten())
+        convolution_input_names.append(bias_tensor_name)
+
+    all_ones = numpy.ones(n_dims, numpy.int8)
+    attrs['dilations'] = all_ones
+    attrs['strides'] = all_ones
+    attrs['kernel_shape'] = all_ones
+    attrs['group'] = 1
+    attrs['auto_pad'] = auto_pad
+
+    intermediate_output_name = scope.get_unique_variable_name('convolution_output')
+    container.add_node('Conv', convolution_input_names,
+                       intermediate_output_name, **attrs)
+    return intermediate_output_name
+
+
 def convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, input_perm_axes,
                             output_perm_axes, weight_perm_axes):
     op = operator.raw_operator
 
+    is_separable_conv = isinstance(op, SeparableConv2D) or \
+                      (StrictVersion(keras.__version__) >= StrictVersion('2.1.3') and isinstance(op, SeparableConv1D))
+
     channels_first = n_dims > 1 and op.data_format == 'channels_first'
 
     # Unless channels_first is the Keras data format, the inputs and weights in Keras v.s. ONNX
@@ -48,10 +82,15 @@ def convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, in
 
     op_type = 'ConvTranspose' if is_transpose else 'Conv'
     convolution_input_names = [adjusted_input_name]
-    attrs = {'name': operator.full_name}
-
     parameters = op.get_weights()
-    assert (len(parameters) == 2 if op.use_bias else 1)
+
+    if is_separable_conv:
+        attrs = {'name': operator.full_name + '0'}
+        assert (len(parameters) == 3 if op.use_bias else 2)
+    else:
+        attrs = {'name': operator.full_name}
+        assert (len(parameters) == 2 if op.use_bias else 1)
+
     weight_params = parameters[0]
 
     input_channels, output_channels = weight_params.shape[-2:]
@@ -68,6 +107,11 @@ def convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, in
         new_shape = shape[:2] + (1, shape[2] * shape[3])
         weight_params = numpy.reshape(weight_params, new_shape)
         weight_params = weight_params.transpose(weight_perm_axes)
+    elif is_separable_conv:
+        group = weight_params.shape[-2]
+        shape = weight_params.shape
+        new_shape = shape[:-2] + (1, shape[-2] * shape[-1])
+        weight_params = numpy.reshape(weight_params, new_shape).transpose(weight_perm_axes)
     else:
         weight_params = weight_params.transpose(weight_perm_axes)
         group = 1
@@ -77,7 +121,7 @@ def convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, in
                               weight_params.shape, weight_params.flatten())
     convolution_input_names.append(weight_tensor_name)
 
-    if len(parameters) == 2:
+    if len(parameters) == 2 and not is_separable_conv:
         bias_tensor_name = scope.get_unique_variable_name('B')
         container.add_initializer(bias_tensor_name, onnx_proto.TensorProto.FLOAT,
                                   parameters[1].shape, parameters[1].flatten())
@@ -114,6 +158,10 @@ def convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, in
     container.add_node(op_type, convolution_input_names,
                        intermediate_output_name, **attrs)
 
+    if is_separable_conv:
+        intermediate_output_name = process_separable_conv_2nd(scope, operator, container, [intermediate_output_name], n_dims,
+                                   weight_perm_axes, parameters, attrs['auto_pad'])
+
     # The construction of convolution is done. Now, we create an activation operator to apply the activation specified
     # in this Keras layer.
     apply_activation_function = _activation_map[op.activation]
@@ -166,10 +214,23 @@ def convert_keras_conv_transpose_3d(scope, operator, container):
     convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, input_perm, output_perm, weight_perm)
 
 
+def convert_keras_separable_conv1d(scope, operator, container):
+    is_transpose, n_dims, input_perm, output_perm, weight_perm = get_converter_config(1, False)
+    convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, input_perm, output_perm, weight_perm)
+
+
+def convert_keras_separable_conv2d(scope, operator, container):
+    is_transpose, n_dims, input_perm, output_perm, weight_perm = get_converter_config(2, False)
+    convert_keras_conv_core(scope, operator, container, is_transpose, n_dims, input_perm, output_perm, weight_perm)
+
+
 register_converter(Conv1D, convert_keras_conv1d)
 register_converter(Conv2D, convert_keras_conv2d)
 register_converter(Conv3D, convert_keras_conv3d)
 register_converter(Conv2DTranspose, convert_keras_conv_transpose_2d)
 register_converter(Conv3DTranspose, convert_keras_conv_transpose_3d)
 if StrictVersion(keras.__version__) >= StrictVersion('2.1.5'):
     register_converter(DepthwiseConv2D, convert_keras_depthwise_conv_2d)
+register_converter(SeparableConv2D, convert_keras_separable_conv2d)
+if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+    register_converter(SeparableConv1D, convert_keras_separable_conv1d)

onnxmltools/convert/keras/shape_calculators/Conv.py

@@ -7,9 +7,11 @@
 import keras
 from distutils.version import StrictVersion
 import numbers
-from keras.layers import Conv1D, Conv2D, Conv3D, Conv2DTranspose, Conv3DTranspose, RepeatVector
+from keras.layers import Conv1D, Conv2D, Conv3D, Conv2DTranspose, Conv3DTranspose, RepeatVector, SeparableConv2D
 if StrictVersion(keras.__version__) >= StrictVersion('2.1.5'):
     from keras.layers import DepthwiseConv2D
+if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+    from keras.layers import SeparableConv1D
 from ...common._registration import register_shape_calculator
 
 
@@ -36,5 +38,8 @@ def calculate_keras_depthwise_conv_output_shapes(operator):
 register_shape_calculator(Conv3D, calculate_keras_conv_output_shapes)
 register_shape_calculator(Conv2DTranspose, calculate_keras_conv_output_shapes)
 register_shape_calculator(Conv3DTranspose, calculate_keras_conv_output_shapes)
+register_shape_calculator(SeparableConv2D, calculate_keras_conv_output_shapes)
 if StrictVersion(keras.__version__) >= StrictVersion('2.1.5'):
     register_shape_calculator(DepthwiseConv2D, calculate_keras_depthwise_conv_output_shapes)
+if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+    register_shape_calculator(SeparableConv1D, calculate_keras_conv_output_shapes)

tests/end2end/test_single_operator_with_cntk_backend.py

@@ -9,14 +9,17 @@
 import onnxmltools
 import coremltools
 import numpy as np
+import keras
 from distutils.version import StrictVersion
 from onnxmltools.proto import onnx
 from onnxmltools.utils.tests_dl_helper import evaluate_deep_model, create_tensor,\
     find_inference_engine, rt_onnxruntime, rt_caffe2, rt_cntk
 from keras.models import Sequential, Model
 from keras.layers import Input, Dense, Conv2D, MaxPooling2D, AveragePooling2D, Conv2DTranspose, \
     Dot, Embedding, BatchNormalization, GRU, Activation, PReLU, LeakyReLU, ThresholdedReLU, Maximum, \
-    Add, Average, Multiply, Concatenate, UpSampling2D, Flatten, RepeatVector, Reshape, Dropout
+    Add, Average, Multiply, Concatenate, UpSampling2D, Flatten, RepeatVector, Reshape, Dropout, SeparableConv2D
+if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+    from keras.layers import SeparableConv1D
 from keras.initializers import RandomUniform
 
 
@@ -422,6 +425,24 @@ def test_recursive_and_shared_model(self):
         # coremltools can't convert this kind of model.
         self._test_one_to_one_operator_keras(model, [x, 2 * x])
 
+    if StrictVersion(keras.__version__) >= StrictVersion('2.1.3'):
+        @unittest.skipIf(find_inference_engine() == rt_cntk, 'Skip because CNTK is not able to evaluate this model')
+        def test_separable_convolution(self):
+            N, C, H, W = 2, 3, 5, 5
+            x = np.random.rand(N, H, W, C).astype(np.float32, copy=False)
+            model = Sequential()
+            model.add(SeparableConv2D(filters=10, kernel_size=(1, 2), strides=(1, 1), padding='valid', input_shape=(H, W, C),
+                             data_format='channels_last', depth_multiplier=4))
+            model.add(MaxPooling2D((2, 2), strides=(2, 2), data_format='channels_last'))
+            model.compile(optimizer='sgd', loss='mse')
+            self._test_one_to_one_operator_keras(model, x)
+
+            x = np.random.rand(N, H, C).astype(np.float32, copy=False)
+            model = Sequential()
+            model.add(SeparableConv1D(filters=10, kernel_size=2, strides=1, padding='valid', input_shape=(H, C),
+                             data_format='channels_last'))
+            model.compile(optimizer='sgd', loss='mse')
+            self._test_one_to_one_operator_keras(model, x)
 
 if __name__ == "__main__":
     unittest.main()