Add a python operator for DepthwiseConvolution

smaug/python/ops/nn_ops.py

@@ -74,6 +74,66 @@ def compute_output_dim(input_dim, weight_dim, stride, padding):
       output_tensors_dims=[output_tensor_dims],
       output_tensor_layout=output_layout, params=params)[0]
 
+def depthwise_convolution(
+    input_tensor, filter_tensor, stride, padding, activation=None,
+    activation_params=None, name="depthwise_conv"):
+  """Compute a 3D depthwise Convolution given 4D `input_tensor` and `filter_tensor`.
+
+  Args:
+    input_tensor: A 4D `Tensor`.
+    filter_tensor: A 4D `Tensor`.
+    stride: A list of two integers: [row_stride, col_stride].
+    padding: A string from: `same`, `valid`. The zero padding options.
+    activation: A string representing the activation function (optional).
+    activation_params: kwargs for the activation function (optional).
+    name: Operator name (optional).
+  """
+  def compute_output_dim(input_dim, weight_dim, stride, padding):
+    pad = 0
+    if to_padding_type(padding) == types_pb2.SamePadding:
+      pad = weight_dim - 1
+    return (input_dim - weight_dim + pad) // stride + 1
+
+  input_tensor, filter_tensor = array_ops.check_and_add_layout_transform(
+      name=name, op=types_pb2.ConvolutionDepthwise,
+      input_tensors=[input_tensor, filter_tensor])
+
+  row_idx = 2 if input_tensor.shape.layout == types_pb2.NCHW else 1
+  col_idx = 3 if input_tensor.shape.layout == types_pb2.NCHW else 2
+  chan_idx = 1 if input_tensor.shape.layout == types_pb2.NCHW else 3
+  assert input_tensor.dims(chan_idx) == filter_tensor.dims(chan_idx), (
+      "The weights must have the same number of channels as the inputs.")
+  output_rows = compute_output_dim(input_tensor.shape.dims[row_idx],
+                                   filter_tensor.shape.dims[row_idx], stride[0],
+                                   padding)
+  output_cols = compute_output_dim(input_tensor.shape.dims[col_idx],
+                                   filter_tensor.shape.dims[col_idx], stride[1],
+                                   padding)
+  output_layout = input_tensor.shape.layout
+  if output_layout == types_pb2.NCHW:
+    output_tensor_dims = [
+        input_tensor.shape.dims[0], input_tensor.shape.dims[chan_idx], output_rows,
+        output_cols
+    ]
+  elif output_layout == types_pb2.NHWC:
+    output_tensor_dims = [
+        input_tensor.shape.dims[0], output_rows, output_cols,
+        input_tensor.shape.dims[chan_idx]
+    ]
+  else:
+    assert False, "Unsupported output layout!"
+  params = node_pb2.Params()
+  params.conv_params.padding = to_padding_type(padding)
+  params.conv_params.stride.extend(stride)
+  if activation is not None:
+    params.act_params.CopyFrom(
+        activation_ops.to_proto(activation, activation_params))
+  return common.add_node(
+      name=name, op=types_pb2.ConvolutionDepthwise,
+      input_tensors=[input_tensor, filter_tensor],
+      output_tensors_dims=[output_tensor_dims],
+      output_tensor_layout=output_layout, params=params)[0]
+
 def batch_norm(
     input_tensor, mean_tensor, var_tensor, gamma_tensor, beta_tensor,
     activation=None, activation_params=None, name="batch_norm"):

smaug/python/ops/ops_test.py

@@ -62,6 +62,9 @@ def build_test_sequential_graph(self, backend):
       filter_tensor1 = Tensor(
           data_layout=types_pb2.NCHW,
           tensor_data=np.random.rand(64, 64, 3, 3).astype(np_dtype))
+      filter_tensor2 = Tensor(
+          data_layout=types_pb2.NCHW,
+          tensor_data=np.random.rand(1, 1, 3, 3).astype(np_dtype))
       weight_tensor0 = Tensor(
           data_layout=types_pb2.NC,
           tensor_data=np.random.rand(254, 12544).astype(np_dtype))
@@ -104,6 +107,8 @@ def build_test_sequential_graph(self, backend):
       out0, out1, out2, out3 = array_ops.unstack(out, 1, "unstack")
       out0 = array_ops.reshape(out0, [1, 1, 8, 10], types_pb2.NCHW, "reshape")
       out0 = array_ops.padding(out0, [0, 0, 0, 0, 1, 1, 1, 1], "padding")
+      out0 = nn_ops.depthwise_convolution(
+          out0, filter_tensor2, stride=[1, 1], padding="same", name="depthwise_conv0")
 
     self.test_graph, _ = graph.to_proto()
     self.backend = backend
@@ -466,6 +471,35 @@ def test_unstack_op(self):
       self.assertEqual(node.output_tensors[0].shape.layout, types_pb2.NC)
       self.assertEqual(node.output_tensors[0].shape.alignment, self.alignment)
 
+  def test_depthwise_convolution_op(self):
+    expected_weight_layout = global_vars.backend_layouts[self.backend][
+        types_pb2.ConvolutionDepthwise].input_layoutsets[1].layouts
+    expected_output_layout = global_vars.backend_layouts[self.backend][
+        types_pb2.ConvolutionDepthwise].output_layoutset.layouts
+    node = self.get_node("depthwise_conv0")
+    self.assertEqual(node.op, types_pb2.ConvolutionDepthwise)
+    self.assertEqual(len(node.input_tensors), 2)
+    self.assertEqual(len(node.output_tensors), 1)
+    # Parameters
+    self.assertEqual(node.params.conv_params.padding, types_pb2.SamePadding)
+    self.assertEqual(node.params.conv_params.stride, [1, 1])
+    # Weight tensor
+    self.assertEqual(node.input_tensors[1].data_type, self.expected_dtype)
+    self.assertEqualDims(node.input_tensors[1].shape.dims,
+                         node.input_tensors[1].shape.layout, [1, 1, 3, 3],
+                         types_pb2.NCHW)
+    self.assertEqual(node.input_tensors[1].shape.layout, expected_weight_layout)
+    self.assertEqual(node.input_tensors[1].shape.alignment, self.alignment)
+    # Output tensor
+    self.assertEqual(node.output_tensors[0].name, "depthwise_conv0/output0")
+    self.assertEqual(node.output_tensors[0].data_type, self.expected_dtype)
+    self.assertEqualDims(node.output_tensors[0].shape.dims,
+                         node.output_tensors[0].shape.layout, [1, 1, 10, 12],
+                         types_pb2.NCHW)
+    self.assertEqual(node.output_tensors[0].shape.layout,
+                     expected_output_layout)
+    self.assertEqual(node.output_tensors[0].shape.alignment, self.alignment)
+
 class ResidualGraphTest(OperatorTest):
   """Common tests for the residual graph."""