From 9e1fcc863bb5e6fb75b3fb1ad054ca3ca51aad99 Mon Sep 17 00:00:00 2001 From: Yida Wang Date: Mon, 1 Jul 2019 15:09:50 -0700 Subject: [PATCH] [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 | 6 +++ src/relay/pass/mac_count.cc | 45 +++++++++++++++++++++-- tests/python/relay/test_pass_mac_count.py | 31 +++++++++++++++- 3 files changed, 77 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 7bce9dd3c5b9..1de86173040d 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/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. diff --git a/src/relay/pass/mac_count.cc b/src/relay/pass/mac_count.cc index 3d77fabe6fe9..ce70eb051214 100644 --- a/src/relay/pass/mac_count.cc +++ b/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 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(); + const auto* data_type = args[0]->checked_type().as(); + Array 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(data_shape[C_ind].as()->value); + if (c_ind != -1) + input_channel *= static_cast(data_shape[c_ind].as()->value); + Array 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(); + Array 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 data_shape = data_type->shape; Array 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(data_shape[0].as()->value); int64_t d2 = static_cast(data_shape[1].as()->value); int64_t d3 = static_cast(weight_shape[0].as()->value); int64_t d4 = static_cast(weight_shape[1].as()->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", ConvMacCount); +RELAY_REGISTER_OP("nn.conv2d_transpose") +.set_attr("FMacCount", Conv2dTransposeMacCount); + RELAY_REGISTER_OP("nn.dense") .set_attr("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_; diff --git a/tests/python/relay/test_pass_mac_count.py b/tests/python/relay/test_pass_mac_count.py index 98ba1ad6325d..a7739a644473 100644 --- a/tests/python/relay/test_pass_mac_count.py +++ b/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()