Skip to content

Commit

Permalink
Add count_include_pad support (#498)
Browse files Browse the repository at this point in the history 
* update tvm submodule

* Add count_include_pad support to avg_pool
  • Loading branch information
@tqchen
Tatsuya Nishiyama authored and tqchen committed May 29, 2018
1 parent 86daf79 commit 918dff2
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 28 deletions.
33 changes: 31 additions & 2 deletions nnvm/include/nnvm/top/nn.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}))
Expand All @@ -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;

Expand Down
16 changes: 8 additions & 8 deletions nnvm/src/compiler/fold_scale_axis.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,45 +354,45 @@ 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];
}
return false;
}

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>);



Expand Down
40 changes: 22 additions & 18 deletions nnvm/src/top/nn/pooling.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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.
Expand All @@ -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")))
Expand All @@ -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)
Expand Down
36 changes: 36 additions & 0 deletions nnvm/tests/python/compiler/test_top_level2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -201,6 +236,7 @@ def test_upsampling():
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()

0 comments on commit 918dff2

Please sign in to comment.