<< "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";
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;
}
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);
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_;
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)
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()