[ANALYSIS] Mac count deconv (#3469)

* add mac count for conv 2d transpose

* add the explanation of missing parameter in docstring

* typo

* fix pylint

python/tvm/relay/op/nn/nn.py

@@ -137,6 +137,12 @@ def conv2d_transpose(data,
     dilation : Tuple[int], optional
         Specifies the dilation rate to be used for dilated convolution.
 
+    channels : int, optional
+        Number of output channels of this convolution.
+
+    kernel_size : tuple of int, optional
+        The spatial of the convolution kernel.
+
     groups : int, optional
         Number of groups for grouped convolution.
 

src/relay/pass/mac_count.cc

@@ -88,11 +88,44 @@ int64_t ConvMacCount(const Call& call_node) {
     << "The dimension of the output tensor in Conv 2D should be 4 or 5.";
   int64_t count = GetCartesianProd(output_tensor) * GetCartesianProd(kernel_size);
   CHECK_EQ(input_channel % conv_2d_attr->groups, 0)
-  << "The number of input channels is not divisble by groups.";
+    << "The number of input channels is not divisble by groups.";
   count *= input_channel/conv_2d_attr->groups;
   return count;
 }
 
+int64_t Conv2dTransposeMacCount(const Call& call_node) {
+  if (!call_node->checked_type_.defined()) {
+    LOG(WARNING) << "The infer type pass should be called before the mac count pass";
+    return 0;
+  }
+  Array<Expr> args = call_node->args;
+  CHECK(args.size() == 2)
+    << "The number of input arguments of a CONV 2D Transpose node should be 2.";
+  const auto* conv_2d_transpose_attr = call_node->attrs.as<Conv2DTransposeAttrs>();
+  const auto* data_type = args[0]->checked_type().as<TensorTypeNode>();
+  Array<IndexExpr> data_shape = data_type->shape;
+  std::string data_layout = conv_2d_transpose_attr->data_layout;
+  int32_t C_ind = Layout(data_layout).IndexOf(LayoutAxis::Get('C'));
+  int32_t c_ind = Layout(data_layout).IndexOf(LayoutAxis::Get('c'));
+  CHECK(C_ind != -1)
+    << "There is no input channel dimension.";
+  int64_t input_channel = static_cast<int64_t>(data_shape[C_ind].as<IntImm>()->value);
+  if (c_ind != -1)
+    input_channel *= static_cast<int64_t>(data_shape[c_ind].as<IntImm>()->value);
+  Array<IndexExpr> kernel_size = conv_2d_transpose_attr->kernel_size;
+  CHECK(kernel_size.size() == 2)
+    << "The dimension of the kernel in Conv 2D Transpose should be 2.";
+  const auto* expr = call_node->checked_type().as<TensorTypeNode>();
+  Array<IndexExpr> output_tensor = expr->shape;
+  CHECK(output_tensor.size() == 4 || output_tensor.size() == 5)
+    << "The dimension of the output tensor in Conv 2D Transpose should be 4 or 5.";
+  int64_t count = GetCartesianProd(output_tensor) * GetCartesianProd(kernel_size);
+  CHECK_EQ(input_channel % conv_2d_transpose_attr->groups, 0)
+    << "The number of input channels is not divisble by groups.";
+  count *= input_channel/conv_2d_transpose_attr->groups;
+  return count;
+}
+
 int64_t DenseMacCount(const Call& call_node) {
   if (!call_node->checked_type_.defined()) {
     LOG(WARNING) << "The infer type pass should be called before the mac count pass";
@@ -106,13 +139,13 @@ int64_t DenseMacCount(const Call& call_node) {
   Array<IndexExpr> data_shape = data_type->shape;
   Array<IndexExpr> weight_shape = weight_type->shape;
   CHECK(data_shape.size() == 2 && weight_shape.size() == 2)
-      << "The dimension of an input tensor to Dense node should be 2.";
+    << "The dimension of an input tensor to Dense node should be 2.";
   int64_t d1 = static_cast<int64_t>(data_shape[0].as<IntImm>()->value);
   int64_t d2 = static_cast<int64_t>(data_shape[1].as<IntImm>()->value);
   int64_t d3 = static_cast<int64_t>(weight_shape[0].as<IntImm>()->value);
   int64_t d4 = static_cast<int64_t>(weight_shape[1].as<IntImm>()->value);
   CHECK(d2 == d4)
-      << "The dimensions of input arguments do not match.";
+    << "The dimensions of input arguments do not match.";
   int64_t count = d1 * d2 * d3;
   return count;
 }
@@ -120,6 +153,9 @@ int64_t DenseMacCount(const Call& call_node) {
 RELAY_REGISTER_OP("nn.conv2d")
 .set_attr<FMacCount>("FMacCount", ConvMacCount);
 
+RELAY_REGISTER_OP("nn.conv2d_transpose")
+.set_attr<FMacCount>("FMacCount", Conv2dTransposeMacCount);
+
 RELAY_REGISTER_OP("nn.dense")
 .set_attr<FMacCount>("FMacCount", DenseMacCount);
 
@@ -129,7 +165,8 @@ class MacCounter : private ExprVisitor {
     count_ = 0;
   }
   static int64_t GetTotalMacNumber(const Expr& expr) {
-    LOG(INFO) << "This pass only counts MACs in direct CONV 2D and Dense ops";
+    LOG(INFO) << "This pass only counts MACs in direct CONV 2D, "
+              << "CONV 2D Transpose and Dense ops";
     MacCounter counter;
     counter(expr);
     return counter.count_;

tests/python/relay/test_pass_mac_count.py

@@ -55,7 +55,7 @@ def test_conv():
         weight,
         channels=output_channel,
         kernel_size=(kh, kw),
-        padding=(1, 1))
+        padding=(h_padding, w_padding))
     func = relay.Function([data, weight],
                             relay.Tuple(tvm.convert([conv2d])))
     func = relay.ir_pass.infer_type(func)
@@ -127,8 +127,37 @@ def test_depthwise_conv2d():
     compute_count = relay.ir_pass.get_total_mac_number(func)
     assert compute_count == 2 * np.prod(dshape) * 3*3
 
+def test_conv_2d_transpose():
+    batch_size = 1
+    input_channel = 3
+    h = 224
+    w = 224
+    output_channel = 64
+    kh = 7
+    kw = 7
+    h_padding = 1
+    w_padding = 1
+    oh = h - h_padding * 2 + kh - 1
+    ow = w - w_padding * 2 + kw - 1
+    dshape = (batch_size, input_channel, h, w)
+    weight = relay.var("weight", shape=(input_channel, output_channel, kh, kw))
+    data = relay.var("data", shape=dshape)
+    conv2d_transpose = relay.nn.conv2d_transpose(
+        data,
+        weight,
+        channels=output_channel,
+        kernel_size=(kh, kw),
+        padding=(h_padding, w_padding))
+    func = relay.Function([data, weight],
+                            relay.Tuple(tvm.convert([conv2d_transpose])))
+    func = relay.ir_pass.infer_type(func)
+    compute_count = relay.ir_pass.get_total_mac_number(func)
+    expect_count = batch_size * input_channel * oh * ow * output_channel * kh * kw
+    assert compute_count == expect_count
+
 if __name__ == "__main__":
     test_conv()
     test_gemm()
     test_simple_network()
     test_depthwise_conv2d()
+    test_conv_2d_transpose()