Add count_include_pad support (#498)

* update tvm submodule

* Add count_include_pad support to avg_pool

nnvm/include/nnvm/top/nn.h

@@ -218,14 +218,14 @@ struct Conv2DTransposeParam : public dmlc::Parameter<Conv2DTransposeParam> {
 };
 
 
-struct Pool2DParam : public dmlc::Parameter<Pool2DParam> {
+struct MaxPool2DParam : public dmlc::Parameter<MaxPool2DParam> {
   TShape pool_size;
   TShape strides;
   TShape padding;
   std::string layout;
   bool ceil_mode;
 
-  DMLC_DECLARE_PARAMETER(Pool2DParam) {
+  DMLC_DECLARE_PARAMETER(MaxPool2DParam) {
     DMLC_DECLARE_FIELD(pool_size)
       .describe("Size of the pooling windows..");
     DMLC_DECLARE_FIELD(strides).set_default(TShape({1, 1}))
@@ -244,6 +244,35 @@ struct Pool2DParam : public dmlc::Parameter<Pool2DParam> {
 };
 
 
+struct AvgPool2DParam : public dmlc::Parameter<AvgPool2DParam> {
+  TShape pool_size;
+  TShape strides;
+  TShape padding;
+  std::string layout;
+  bool ceil_mode;
+  bool count_include_pad;
+
+  DMLC_DECLARE_PARAMETER(AvgPool2DParam) {
+    DMLC_DECLARE_FIELD(pool_size)
+      .describe("Size of the pooling windows..");
+    DMLC_DECLARE_FIELD(strides).set_default(TShape({1, 1}))
+      .describe("Specifies the strides of the convolution.");
+    DMLC_DECLARE_FIELD(padding).set_default(TShape({0, 0}))
+      .describe("If padding is non-zero, then the input is implicitly zero-padded"
+                "on both sides for padding number of points");
+    DMLC_DECLARE_FIELD(layout).set_default("NCHW")
+      .describe("Dimension ordering of data and weight. Can be 'NCHW', 'NHWC', etc."
+                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                "dimensions respectively. Convolution is applied on the 'H' and"
+                "'W' dimensions.");
+    DMLC_DECLARE_FIELD(ceil_mode).set_default(false)
+      .describe("When true, will use ceil instead of floor to compute the output shape.");
+    DMLC_DECLARE_FIELD(count_include_pad).set_default(false)
+      .describe("When true, will include padding to compute the average");
+  }
+};
+
+
 struct GlobalPool2DParam : public dmlc::Parameter<GlobalPool2DParam> {
   std::string layout;
 

nnvm/src/compiler/fold_scale_axis.cc

@@ -354,28 +354,28 @@ NNVM_REGISTER_OP(leaky_relu)
 .set_attr<FScaleAxisForward>("FScaleAxisForward", ReluScaleAxisForward);
 
 // property registration.
+template <typename T>
 bool Pool2DBackward(
     const NodeAttrs& attrs,
     const std::vector<TShape>& in_shape,
     const std::vector<TShape>& out_shape,
     const FoldChainInfo& out_info,
     std::vector<FoldChainInfo>* in_axis) {
-  using top::Pool2DParam;
-  const Pool2DParam& param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const T& param = nnvm::get<T>(attrs.parsed);
   if (out_info.axis == 1 && param.layout == "NCHW") {
     (*in_axis)[0] = out_info;
   }
   return false;
 }
 
+template <typename T>
 bool Pool2DForward(
     const NodeAttrs& attrs,
     const std::vector<TShape>& in_shape,
     const std::vector<TShape>& out_shape,
     std::vector<FoldChainInfo>* in_info,
     FoldChainInfo* out_info) {
-  using top::Pool2DParam;
-  const Pool2DParam& param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const T& param = nnvm::get<T>(attrs.parsed);
   if ((*in_info)[0].axis == 1 && param.layout == "NCHW") {
     *out_info = (*in_info)[0];
   }
@@ -383,16 +383,16 @@ bool Pool2DForward(
 }
 
 NNVM_REGISTER_OP(max_pool2d)
-.set_attr<FScaleAxisBackward>("FScaleAxisBackward", Pool2DBackward);
+.set_attr<FScaleAxisBackward>("FScaleAxisBackward", Pool2DBackward<top::MaxPool2DParam>);
 
 NNVM_REGISTER_OP(avg_pool2d)
-.set_attr<FScaleAxisBackward>("FScaleAxisBackward", Pool2DBackward);
+.set_attr<FScaleAxisBackward>("FScaleAxisBackward", Pool2DBackward<top::AvgPool2DParam>);
 
 NNVM_REGISTER_OP(max_pool2d)
-.set_attr<FScaleAxisForward>("FScaleAxisForward", Pool2DForward);
+.set_attr<FScaleAxisForward>("FScaleAxisForward", Pool2DForward<top::MaxPool2DParam>);
 
 NNVM_REGISTER_OP(avg_pool2d)
-.set_attr<FScaleAxisForward>("FScaleAxisForward", Pool2DForward);
+.set_attr<FScaleAxisForward>("FScaleAxisForward", Pool2DForward<top::AvgPool2DParam>);
 
 
 

nnvm/src/top/nn/pooling.cc

@@ -19,12 +19,13 @@ namespace top {
 using namespace tvm;
 using namespace nnvm::compiler;
 
-DMLC_REGISTER_PARAMETER(Pool2DParam);
+DMLC_REGISTER_PARAMETER(MaxPool2DParam);
 
+template <typename T>
 inline bool Pool2DInferShape(const nnvm::NodeAttrs& attrs,
                              std::vector<TShape>* in_shape,
                              std::vector<TShape>* out_shape) {
-  const Pool2DParam& param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const T& param = nnvm::get<T>(attrs.parsed);
   CHECK_EQ(in_shape->size(), 1U);
   CHECK_EQ(out_shape->size(), 1U);
 
@@ -66,11 +67,12 @@ inline bool Pool2DInferShape(const nnvm::NodeAttrs& attrs,
   return true;
 }
 
+template <typename T>
 inline bool Pool2DCorrectLayout(const NodeAttrs& attrs,
                                 std::vector<Layout> *ilayouts,
                                 const std::vector<Layout> *last_ilayouts,
                                 std::vector<Layout> *olayouts) {
-  const Pool2DParam &param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const T &param = nnvm::get<T>(attrs.parsed);
   CHECK_EQ(ilayouts->size(), 1);
   CHECK_EQ(last_ilayouts->size(), 1);
   CHECK_EQ(olayouts->size(), 1);
@@ -114,18 +116,18 @@ NNVM_REGISTER_OP(max_pool2d)
 
 )code" NNVM_ADD_FILELINE)
 .add_argument("data", "4D Tensor", "Input data.")
-.add_arguments(Pool2DParam::__FIELDS__())
-.set_attr_parser(ParamParser<Pool2DParam>)
-.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<Pool2DParam>)
+.add_arguments(MaxPool2DParam::__FIELDS__())
+.set_attr_parser(ParamParser<MaxPool2DParam>)
+.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<MaxPool2DParam>)
 .set_num_outputs(1)
 .set_num_inputs(1)
-.set_attr<FInferShape>("FInferShape", Pool2DInferShape)
+.set_attr<FInferShape>("FInferShape", Pool2DInferShape<MaxPool2DParam>)
 .set_attr<FInferType>("FInferType", ElemwiseType<1, 1>)
-.set_attr<FCorrectLayout>("FCorrectLayout", Pool2DCorrectLayout)
+.set_attr<FCorrectLayout>("FCorrectLayout", Pool2DCorrectLayout<MaxPool2DParam>)
 .set_attr<FTVMCompute>("FTVMCompute", [](const NodeAttrs& attrs,
                                          const Array<Tensor>& inputs,
                                          const Array<Tensor>& out_info) {
-  const Pool2DParam& param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const MaxPool2DParam& param = nnvm::get<MaxPool2DParam>(attrs.parsed);
   auto pool_size = ShapeToArray(param.pool_size);
   auto strides = ShapeToArray(param.strides);
   auto padding = ShapeToArray(param.padding);
@@ -163,12 +165,13 @@ NNVM_REGISTER_OP(_max_pool2d_grad)
 .add_argument("output", "4D Tensor", "Output data of max_pool2d grad.")
 .set_num_inputs(3)
 .set_num_outputs(1)
-.set_attr_parser(ParamParser<Pool2DParam>)
-.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<Pool2DParam>)
+.set_attr_parser(ParamParser<MaxPool2DParam>)
+.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<MaxPool2DParam>)
 .set_attr<FInferShape>("FInferShape", AssignOutputAttr<TShape, 1, 0>)
 .set_attr<FInferType>("FInferType", ElemwiseType<3, 1>)
 .set_attr<TIsBackward>("TIsBackward", true);
 
+DMLC_REGISTER_PARAMETER(AvgPool2DParam);
 
 NNVM_REGISTER_OP(avg_pool2d)
 .describe(R"code(Average pooling operation for one dimensional data.
@@ -187,20 +190,21 @@ NNVM_REGISTER_OP(avg_pool2d)
 
 )code" NNVM_ADD_FILELINE)
 .add_argument("data", "4D Tensor", "Input data.")
-.add_arguments(Pool2DParam::__FIELDS__())
-.set_attr_parser(ParamParser<Pool2DParam>)
-.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<Pool2DParam>)
-.set_attr<FInferShape>("FInferShape", Pool2DInferShape)
+.add_arguments(AvgPool2DParam::__FIELDS__())
+.set_attr_parser(ParamParser<AvgPool2DParam>)
+.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<AvgPool2DParam>)
+.set_attr<FInferShape>("FInferShape", Pool2DInferShape<AvgPool2DParam>)
 .set_attr<FInferType>("FInferType", ElemwiseType<1, 1>)
-.set_attr<FCorrectLayout>("FCorrectLayout", Pool2DCorrectLayout)
+.set_attr<FCorrectLayout>("FCorrectLayout", Pool2DCorrectLayout<AvgPool2DParam>)
 .set_attr<FTVMCompute>("FTVMCompute", [](const NodeAttrs& attrs,
                                          const Array<Tensor>& inputs,
                                          const Array<Tensor>& out_info) {
-  const Pool2DParam& param = nnvm::get<Pool2DParam>(attrs.parsed);
+  const AvgPool2DParam& param = nnvm::get<AvgPool2DParam>(attrs.parsed);
   auto pool_size = ShapeToArray(param.pool_size);
   auto strides = ShapeToArray(param.strides);
   auto padding = ShapeToArray(param.padding);
   auto ceil_mode = param.ceil_mode;
+  auto count_include_pad = param.count_include_pad;
 
   Layout layout(param.layout);
   CHECK(layout.convertible(Layout("NCHW")))
@@ -214,7 +218,7 @@ NNVM_REGISTER_OP(avg_pool2d)
 
   return Array<Tensor>{
     topi::nn::pool(inputs[0], pool_size, strides, padding,
-                   topi::nn::kAvgPool, ceil_mode, layout.name())};
+                   topi::nn::kAvgPool, ceil_mode, layout.name(), count_include_pad)};
 })
 .set_num_outputs(1)
 .set_num_inputs(1)

nnvm/tests/python/compiler/test_top_level2.py

@@ -141,6 +141,41 @@ def test_avg_pool2d():
         np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
 
 
+def test_avg_pool2d_no_count_pad():
+    kh, kw = (4, 4)
+    sh, sw = (2, 2)
+    ph, pw = (2, 2)
+    
+    x = sym.Variable("x")
+    y = sym.avg_pool2d(x, pool_size=(kh, kw), strides=(sw, sw), padding=(ph, pw),
+                       name="y", count_include_pad=False)
+    dtype = "float32"
+    n = 1
+    (ic, ih, iw) = (3, 28, 28)
+    (oc, oh, ow) = (3, 15, 15)
+
+    a_np = np.random.uniform(low=0.001, size=(n, ic, ih, iw)).astype(dtype)
+    pad_np = np.zeros(shape=(n, ic, ih+2*ph, iw+2*pw)).astype(dtype)
+    no_zero = (range(n), range(ic), (range(ph, ih+ph)), (range(pw, iw+pw)))
+    pad_np[np.ix_(*no_zero)] = a_np
+    b_np = np.zeros(shape=(n, oc, oh, ow)).astype(dtype)
+    
+    for i in range(oh):
+        for j in range(ow):
+            pad_count = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw] > 0, axis=(2,3))
+            b_np[:,:,i,j] = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw],
+                                   axis=(2,3)) / np.maximum(pad_count, 1)
+    b_np = np.maximum(b_np, 0.0)
+    shape_dict = {"x": (n, ic, ih, iw)}
+    for target, ctx in ctx_list():
+        graph, lib, _ = nnvm.compiler.build(y, target, shape_dict)
+        m = graph_runtime.create(graph, lib, ctx)
+        data = tvm.nd.array(a_np)
+        m.run(x=data)
+        out = m.get_output(0, tvm.nd.empty((n, oc, oh, ow), dtype))
+        np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5)
+
+
 def test_global_max_pool2d():
     x = sym.Variable("x")
     y = sym.global_max_pool2d(x, name="y")
@@ -201,6 +236,7 @@ if __name__ == "__main__":
     test_conv2d_transpose()
     test_max_pool2d()
     test_avg_pool2d()
+    test_avg_pool2d_no_count_pad()
     test_global_max_pool2d()
     test_global_avg_pool2d()
     test_upsampling()