Handle conv layer spec when input0 is type SpaceToBatchND

keras2onnx/_builtin.py

@@ -118,43 +118,45 @@ def convert_tf_batch_to_space(scope, operator, container):
                                       starts=[0], ends=[1])
         block_prod = oopb.apply_mul(block_size + block_size,
                                     name=operator.full_name + '_mul_0')
-        padded_block_prod = oopb.apply_pad([block_prod,
-                                            ('const_zero_three', oopb.int64, np.array([0, 3], dtype='int64')),
-                                            ('one', oopb.int64, np.array([1], dtype='int64'))],
-                                           name=operator.full_name + '_pad_0')
+        padded_block_prod = oopb.apply_pad(block_prod,
+                                           name=operator.full_name + '_pad_0',
+                                           pads=[0, 3],
+                                           value=1)
         new_shape_x = oopb.apply_div([shape_x] + padded_block_prod,
                                      name=operator.full_name + '_div')
         concat_new_shape_x = oopb.apply_concat(block_shape + new_shape_x,
                                                name=operator.full_name + '_concat',
                                                axis=0)
-        reshaped_x = oopb.apply_reshape([operator.inputs[0].full_name] + concat_new_shape_x,
-                                        name=operator.full_name + '_reshape_0')
+        reshaped_x = oopb.apply_reshape([operator.inputs[0].full_name],
+                                        name=operator.full_name + '_reshape_0',
+                                        desired_shape=concat_new_shape_x[0])
         transposed_x = oopb.apply_transpose(reshaped_x,
                                             name=operator.full_name + '_transpose_0',
                                             perm=[2, 3, 0, 4, 1, 5])
-        padded_block_shape = oopb.apply_pad([block_shape[0],
-                                             ('const_one_one', oopb.int64, np.array([1, 1], dtype='int64')),
-                                             ('one', oopb.int64, np.array([1], dtype='int64'))],
-                                            name=operator.full_name + '_pad_1')
+        padded_block_shape = oopb.apply_pad(block_shape,
+                                            name = operator.full_name + '_pad_1',
+                                            pads=[1, 1],
+                                            value=1)
         new_shape_x_v2 = oopb.apply_mul(new_shape_x + padded_block_shape,
                                         name=operator.full_name + '_mul_1')
-        reshaped_x_v2 = oopb.apply_reshape(transposed_x + new_shape_x_v2,
-                                           name=operator.full_name + '_reshape_1')
+        reshaped_x_v2 = oopb.apply_reshape(transposed_x,
+                                           name=operator.full_name + '_reshape_1',
+                                           desired_shape=new_shape_x_v2[0])
         transposed_crop = oopb.apply_transpose(crop,
                                                name=operator.full_name + '_transpose_1',
                                                perm=[1, 0])
         slice_crop_starts = oopb.apply_slice(transposed_crop,
                                              name=operator.full_name + '_slice_starts',
                                              starts=[0, 0], ends=[1, 2])
-        reshaped_slice_crop_starts = oopb.apply_reshape(slice_crop_starts +
-                                                        [('const_one_one', oopb.int64, np.array([2], dtype='int64'))],
-                                                        name=operator.full_name + '_reshape_starts')
+        reshaped_slice_crop_starts = oopb.apply_reshape(slice_crop_starts,
+                                                        name=operator.full_name + '_reshape_starts',
+                                                        desired_shape=[2])
         slice_crop_ends = oopb.apply_slice(transposed_crop,
                                            name=operator.full_name + '_slice_ends',
                                            starts=[1, 0], ends=[2, 2])
-        reshaped_slice_crop_ends = oopb.apply_reshape(slice_crop_ends +
-                                                      [('const_two', oopb.int64, np.array([2], dtype='int64'))],
-                                                      name=operator.full_name + '_reshape_ends')
+        reshaped_slice_crop_ends = oopb.apply_reshape(slice_crop_ends,
+                                                      name=operator.full_name + '_reshape_ends',
+                                                      desired_shape=[2])
         sliced_new_shape_x_v2 = oopb.apply_slice(new_shape_x_v2,
                                                  name=operator.full_name + '_slice_3',
                                                  starts=[1], ends=[3])
@@ -224,9 +226,9 @@ def convert_tf_space_to_batch(scope, operator, container):
         transposed_concated_pad_x_v2 = oopb.apply_transpose(concated_pad_x_v2,
                                                             name=operator.full_name + '_transpose_0',
                                                             perm=[1, 0])
-        reshaped_transposed_pad_x = oopb.apply_reshape(transposed_concated_pad_x_v2 +
-                                                       [('_const_eight', oopb.int64, np.array([8], dtype='int64'))],
-                                                       name=operator.full_name + '_reshape_0')
+        reshaped_transposed_pad_x = oopb.apply_reshape(transposed_concated_pad_x_v2,
+                                                       name=operator.full_name + '_reshape_0',
+                                                       desired_shape=[8])
         padded_input_x = oopb.apply_pad(operator.inputs[0].full_name,
                                         name=operator.full_name + '_pad_1',
                                         pads=reshaped_transposed_pad_x)
@@ -260,8 +262,9 @@ def convert_tf_space_to_batch(scope, operator, container):
         new_shape_x_v2 = oopb.apply_concat(new_first_row_new_shape_x + new_second_row_new_shape_x,
                                            name=operator.full_name + '_concat_shape',
                                            axis=0)
-        new_x = oopb.apply_reshape(padded_input_x + new_shape_x_v2,
-                                   name=operator.full_name + '_reshape_new_x')
+        new_x = oopb.apply_reshape(padded_input_x[0],
+                                   name=operator.full_name + '_reshape_new_x',
+                                   desired_shape=new_shape_x_v2[0])
         transposed_new_x = oopb.apply_transpose(new_x,
                                                 name=operator.full_name + '_transpose_new',
                                                 perm=[2, 4, 0, 1, 3, 5])

keras2onnx/ke2onnx/layer_spec.py

@@ -52,7 +52,10 @@ def _conv_layer_spec_outputs(layer, node):
         return node.name[:ri + 1] + node_act
     else:
         if not layer.use_bias:
-            return node.name
+            if node.inputs[0].op.type == 'SpaceToBatchND':
+                return node.name + '/BatchToSpaceND'
+            else:
+                return node.name
         else:
             ri = node.name.rindex('/')
             return node.name[:ri + 1] + 'BiasAdd'

tests/test_layers.py

@@ -192,12 +192,13 @@ def my_func_2(x):
         assert runner('onnx_concat', onnx_model, [data1, data2], expected)
 
 
-def test_depthwise_conv2d(runner):
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_depthwise_conv2d(runner, use_bias):
     model = Sequential()
     model.add(InputLayer(input_shape=(8, 8, 2)))
     model.add(keras.layers.DepthwiseConv2D(
         kernel_size=(3, 3), strides=(1, 1), padding="VALID",
-        data_format='channels_last'))
+        data_format='channels_last', use_bias=use_bias))
     onnx_model = keras2onnx.convert_keras(model, 'test_depthwise_conv2d')
     data = np.random.rand(3, 8, 8, 2).astype(np.float32)
     expected = model.predict(data)