Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
Add NHWC layout support to Pooling (cpu, gpu cuda, gpu cuDNN) (#13749)
Browse files Browse the repository at this point in the history 
* Adds layout support: mx.sym.Pooling(..., layout='NHWC',...) with tests.

* Docs changes

* Trigger

* Skip NHWC pooling tests on non-cuDNN platforms

* Fix pylint NHWC pooling

* Fixes from review

* Add CuDNNPoolingOp::Supports() in place of Forward()/Backward() bool return.

* Add layout support to cpu implementation of Pooling, with tests.

* Fix cpplint.

* Fix bug in cpu nhwc impl.

* Add MXNet CUDA pooling in NWC, NHWC and NDHWC.  Turn on 3D cuDNN pooling.  Tests.

* Add PoolingParam::GetLayout() for better default layout handling.

* Fix cpplint.

* Throw exception for quantization pooling not NCHW.

* Expand nhwc pooling test coverage.

* SupportMKLDNNPooling() to examine layout param.

* Compare 'std' and 'v1' pooling versions only when op definitions permit.

* Add pooling test diagnostic output.

* Fix syntax.

* Fix pooling FInplaceOption so it can be shared by all implementations.

* Add missing param definition.

* Fix #if logic.

* Temp switch to DickJC123/mshadow: shows effect of half round-to-nearest on cpu.

* Move back to dmlc/mshadow.git, now with float->half rounding.

* Avoid underflow of lp pooling calc for dtype=float16.

* Remove redundant pooling test.

* Minor variable naming fixes.

* Modify FInplaceOption handling per reviewer comments.  Expand testing.

* Correct gluon Pooling layout param description.

* Correct Symbol Pooling description.

* Use 'CHECK(x)' rather than 'if (x) LOG(FATAL)'.

* Empty commit to trigger CI.
  • Loading branch information
@DickJC123 @eric-haibin-lin
DickJC123 authored and eric-haibin-lin committed Feb 16, 2019
1 parent 059f49f commit 5adb6fc
Show file tree
Hide file tree
Showing 12 changed files with 2,024 additions and 754 deletions.
93 changes: 55 additions & 38 deletions python/mxnet/gluon/nn/conv_layers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -673,7 +673,7 @@ def __init__(self, channels, kernel_size, strides=(1, 1, 1), padding=(0, 0, 0),
class _Pooling(HybridBlock):
"""Abstract class for different pooling layers."""
def __init__(self, pool_size, strides, padding, ceil_mode, global_pool,
pool_type, count_include_pad=None, **kwargs):
pool_type, layout, count_include_pad=None, **kwargs):
super(_Pooling, self).__init__(**kwargs)
if strides is None:
strides = pool_size
Expand All @@ -684,6 +684,7 @@ def __init__(self, pool_size, strides, padding, ceil_mode, global_pool,
self._kwargs = {
'kernel': pool_size, 'stride': strides, 'pad': padding,
'global_pool': global_pool, 'pool_type': pool_type,
'layout': layout,
'pooling_convention': 'full' if ceil_mode else 'valid'}
if count_include_pad is not None:
self._kwargs['count_include_pad'] = count_include_pad
Expand All @@ -695,7 +696,8 @@ def hybrid_forward(self, F, x):
return F.Pooling(x, name='fwd', **self._kwargs)

def __repr__(self):
s = '{name}(size={kernel}, stride={stride}, padding={pad}, ceil_mode={ceil_mode})'
s = '{name}(size={kernel}, stride={stride}, padding={pad}, ceil_mode={ceil_mode}'
s += ', global_pool={global_pool}, pool_type={pool_type}, layout={layout})'
return s.format(name=self.__class__.__name__,
ceil_mode=self._kwargs['pooling_convention'] == 'full',
**self._kwargs)
Expand All @@ -716,7 +718,7 @@ class MaxPool1D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCW'
Dimension ordering of data and weight. Only supports 'NCW' layout for now.
Dimension ordering of data and out ('NCW' or 'NWC').
'N', 'C', 'W' stands for batch, channel, and width (time) dimensions
respectively. Pooling is applied on the W dimension.
ceil_mode : bool, default False
Expand All @@ -738,12 +740,13 @@ class MaxPool1D(_Pooling):
"""
def __init__(self, pool_size=2, strides=None, padding=0, layout='NCW',
ceil_mode=False, **kwargs):
assert layout == 'NCW', "Only supports 'NCW' layout for now"
assert layout in ('NCW', 'NWC'),\
"Only NCW and NWC layouts are valid for 1D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)
assert len(pool_size) == 1, "pool_size must be a number or a list of 1 ints"
super(MaxPool1D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'max', **kwargs)
pool_size, strides, padding, ceil_mode, False, 'max', layout, **kwargs)


class MaxPool2D(_Pooling):
Expand All @@ -761,7 +764,7 @@ class MaxPool2D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCHW'
Dimension ordering of data and weight. Only supports 'NCHW' layout for now.
Dimension ordering of data and out ('NCHW' or 'NHWC').
'N', 'C', 'H', 'W' stands for batch, channel, height, and width
dimensions respectively. padding is applied on 'H' and 'W' dimension.
ceil_mode : bool, default False
Expand All @@ -786,12 +789,13 @@ class MaxPool2D(_Pooling):
"""
def __init__(self, pool_size=(2, 2), strides=None, padding=0, layout='NCHW',
ceil_mode=False, **kwargs):
assert layout == 'NCHW', "Only supports 'NCHW' layout for now"
assert layout in ('NCHW', 'NHWC'),\
"Only NCHW and NHWC layouts are valid for 2D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)*2
assert len(pool_size) == 2, "pool_size must be a number or a list of 2 ints"
super(MaxPool2D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'max', **kwargs)
pool_size, strides, padding, ceil_mode, False, 'max', layout, **kwargs)


class MaxPool3D(_Pooling):
Expand All @@ -809,7 +813,7 @@ class MaxPool3D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCDHW'
Dimension ordering of data and weight. Only supports 'NCDHW' layout for now.
Dimension ordering of data and out ('NCDHW' or 'NDHWC').
'N', 'C', 'H', 'W', 'D' stands for batch, channel, height, width and
depth dimensions respectively. padding is applied on 'D', 'H' and 'W'
dimension.
Expand All @@ -836,12 +840,13 @@ class MaxPool3D(_Pooling):
"""
def __init__(self, pool_size=(2, 2, 2), strides=None, padding=0,
ceil_mode=False, layout='NCDHW', **kwargs):
assert layout == 'NCDHW', "Only supports 'NCDHW' layout for now"
assert layout in ('NCDHW', 'NDHWC'),\
"Only NCDHW and NDHWC layouts are valid for 3D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)*3
assert len(pool_size) == 3, "pool_size must be a number or a list of 3 ints"
super(MaxPool3D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'max', **kwargs)
pool_size, strides, padding, ceil_mode, False, 'max', layout, **kwargs)


class AvgPool1D(_Pooling):
Expand All @@ -858,7 +863,7 @@ class AvgPool1D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCW'
Dimension ordering of data and weight. Only supports 'NCW' layout for now.
Dimension ordering of data and out ('NCW' or 'NWC').
'N', 'C', 'W' stands for batch, channel, and width (time) dimensions
respectively. padding is applied on 'W' dimension.
ceil_mode : bool, default False
Expand All @@ -882,12 +887,14 @@ class AvgPool1D(_Pooling):
"""
def __init__(self, pool_size=2, strides=None, padding=0, layout='NCW',
ceil_mode=False, count_include_pad=True, **kwargs):
assert layout == 'NCW', "Only supports 'NCW' layout for now"
assert layout in ('NCW', 'NWC'),\
"Only NCW and NWC layouts are valid for 1D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)
assert len(pool_size) == 1, "pool_size must be a number or a list of 1 ints"
super(AvgPool1D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'avg', count_include_pad, **kwargs)
pool_size, strides, padding, ceil_mode, False, 'avg', layout, count_include_pad,
**kwargs)


class AvgPool2D(_Pooling):
Expand All @@ -904,7 +911,7 @@ class AvgPool2D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCHW'
Dimension ordering of data and weight. Only supports 'NCHW' layout for now.
Dimension ordering of data and out ('NCHW' or 'NHWC').
'N', 'C', 'H', 'W' stands for batch, channel, height, and width
dimensions respectively. padding is applied on 'H' and 'W' dimension.
ceil_mode : bool, default False
Expand All @@ -931,12 +938,14 @@ class AvgPool2D(_Pooling):
"""
def __init__(self, pool_size=(2, 2), strides=None, padding=0,
ceil_mode=False, layout='NCHW', count_include_pad=True, **kwargs):
assert layout == 'NCHW', "Only supports 'NCHW' layout for now"
assert layout in ('NCHW', 'NHWC'),\
"Only NCHW and NHWC layouts are valid for 2D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)*2
assert len(pool_size) == 2, "pool_size must be a number or a list of 2 ints"
super(AvgPool2D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'avg', count_include_pad, **kwargs)
pool_size, strides, padding, ceil_mode, False, 'avg', layout, count_include_pad,
**kwargs)


class AvgPool3D(_Pooling):
Expand All @@ -953,7 +962,7 @@ class AvgPool3D(_Pooling):
If padding is non-zero, then the input is implicitly
zero-padded on both sides for padding number of points.
layout : str, default 'NCDHW'
Dimension ordering of data and weight. Can be 'NCDHW', 'NDHWC', etc.
Dimension ordering of data and out ('NCDHW' or 'NDHWC').
'N', 'C', 'H', 'W', 'D' stands for batch, channel, height, width and
depth dimensions respectively. padding is applied on 'D', 'H' and 'W'
dimension.
Expand Down Expand Up @@ -982,12 +991,14 @@ class AvgPool3D(_Pooling):
"""
def __init__(self, pool_size=(2, 2, 2), strides=None, padding=0,
ceil_mode=False, layout='NCDHW', count_include_pad=True, **kwargs):
assert layout == 'NCDHW', "Only supports 'NCDHW' layout for now"
assert layout in ('NCDHW', 'NDHWC'),\
"Only NCDHW and NDHWC layouts are valid for 3D Pooling"
if isinstance(pool_size, numeric_types):
pool_size = (pool_size,)*3
assert len(pool_size) == 3, "pool_size must be a number or a list of 3 ints"
super(AvgPool3D, self).__init__(
pool_size, strides, padding, ceil_mode, False, 'avg', count_include_pad, **kwargs)
pool_size, strides, padding, ceil_mode, False, 'avg', layout, count_include_pad,
**kwargs)


class GlobalMaxPool1D(_Pooling):
Expand All @@ -997,7 +1008,7 @@ class GlobalMaxPool1D(_Pooling):
Parameters
----------
layout : str, default 'NCW'
Dimension ordering of data and weight. Only supports 'NCW' layout for now.
Dimension ordering of data and out ('NCW' or 'NWC').
'N', 'C', 'W' stands for batch, channel, and width (time) dimensions
respectively. Pooling is applied on the W dimension.
Expand All @@ -1011,9 +1022,10 @@ class GlobalMaxPool1D(_Pooling):
when `layout` is `NCW`.
"""
def __init__(self, layout='NCW', **kwargs):
assert layout == 'NCW', "Only supports 'NCW' layout for now"
assert layout in ('NCW', 'NWC'),\
"Only NCW and NWC layouts are valid for 1D Pooling"
super(GlobalMaxPool1D, self).__init__(
(1,), None, 0, True, True, 'max', **kwargs)
(1,), None, 0, True, True, 'max', layout, **kwargs)


class GlobalMaxPool2D(_Pooling):
Expand All @@ -1023,7 +1035,7 @@ class GlobalMaxPool2D(_Pooling):
Parameters
----------
layout : str, default 'NCHW'
Dimension ordering of data and weight. Only supports 'NCHW' layout for now.
Dimension ordering of data and out ('NCHW' or 'NHWC').
'N', 'C', 'H', 'W' stands for batch, channel, height, and width
dimensions respectively. padding is applied on 'H' and 'W' dimension.
Expand All @@ -1038,9 +1050,10 @@ class GlobalMaxPool2D(_Pooling):
`(batch_size, channels, 1, 1)` when `layout` is `NCHW`.
"""
def __init__(self, layout='NCHW', **kwargs):
assert layout == 'NCHW', "Only supports 'NCHW' layout for now"
assert layout in ('NCHW', 'NHWC'),\
"Only NCHW and NHWC layouts are valid for 2D Pooling"
super(GlobalMaxPool2D, self).__init__(
(1, 1), None, 0, True, True, 'max', **kwargs)
(1, 1), None, 0, True, True, 'max', layout, **kwargs)


class GlobalMaxPool3D(_Pooling):
Expand All @@ -1050,7 +1063,7 @@ class GlobalMaxPool3D(_Pooling):
Parameters
----------
layout : str, default 'NCDHW'
Dimension ordering of data and weight. Only supports 'NCDHW' layout for now.
Dimension ordering of data and out ('NCDHW' or 'NDHWC').
'N', 'C', 'H', 'W', 'D' stands for batch, channel, height, width and
depth dimensions respectively. padding is applied on 'D', 'H' and 'W'
dimension.
Expand All @@ -1066,9 +1079,10 @@ class GlobalMaxPool3D(_Pooling):
`(batch_size, channels, 1, 1, 1)` when `layout` is `NCDHW`.
"""
def __init__(self, layout='NCDHW', **kwargs):
assert layout == 'NCDHW', "Only supports 'NCDHW' layout for now"
assert layout in ('NCDHW', 'NDHWC'),\
"Only NCDHW and NDHWC layouts are valid for 3D Pooling"
super(GlobalMaxPool3D, self).__init__(
(1, 1, 1), None, 0, True, True, 'max', **kwargs)
(1, 1, 1), None, 0, True, True, 'max', layout, **kwargs)


class GlobalAvgPool1D(_Pooling):
Expand All @@ -1077,7 +1091,7 @@ class GlobalAvgPool1D(_Pooling):
Parameters
----------
layout : str, default 'NCW'
Dimension ordering of data and weight. Only supports 'NCW' layout for now.
Dimension ordering of data and out ('NCW' or 'NWC').
'N', 'C', 'W' stands for batch, channel, and width (time) dimensions
respectively. padding is applied on 'W' dimension.
Expand All @@ -1090,9 +1104,10 @@ class GlobalAvgPool1D(_Pooling):
- **out**: 3D output tensor with shape `(batch_size, channels, 1)`.
"""
def __init__(self, layout='NCW', **kwargs):
assert layout == 'NCW', "Only supports 'NCW' layout for now"
assert layout in ('NCW', 'NWC'),\
"Only NCW and NWC layouts are valid for 1D Pooling"
super(GlobalAvgPool1D, self).__init__(
(1,), None, 0, True, True, 'avg', **kwargs)
(1,), None, 0, True, True, 'avg', layout, **kwargs)


class GlobalAvgPool2D(_Pooling):
Expand All @@ -1101,7 +1116,7 @@ class GlobalAvgPool2D(_Pooling):
Parameters
----------
layout : str, default 'NCHW'
Dimension ordering of data and weight. Only supports 'NCHW' layout for now.
Dimension ordering of data and out ('NCHW' or 'NHWC').
'N', 'C', 'H', 'W' stands for batch, channel, height, and width
dimensions respectively.
Expand All @@ -1116,9 +1131,10 @@ class GlobalAvgPool2D(_Pooling):
`(batch_size, channels, 1, 1)` when `layout` is `NCHW`.
"""
def __init__(self, layout='NCHW', **kwargs):
assert layout == 'NCHW', "Only supports 'NCHW' layout for now"
assert layout in ('NCHW', 'NHWC'),\
"Only NCHW and NHWC layouts are valid for 2D Pooling"
super(GlobalAvgPool2D, self).__init__(
(1, 1), None, 0, True, True, 'avg', **kwargs)
(1, 1), None, 0, True, True, 'avg', layout, **kwargs)


class GlobalAvgPool3D(_Pooling):
Expand All @@ -1127,7 +1143,7 @@ class GlobalAvgPool3D(_Pooling):
Parameters
----------
layout : str, default 'NCDHW'
Dimension ordering of data and weight. Can be 'NCDHW', 'NDHWC', etc.
Dimension ordering of data and out ('NCDHW' or 'NDHWC').
'N', 'C', 'H', 'W', 'D' stands for batch, channel, height, width and
depth dimensions respectively. padding is applied on 'D', 'H' and 'W'
dimension.
Expand All @@ -1143,9 +1159,10 @@ class GlobalAvgPool3D(_Pooling):
`(batch_size, channels, 1, 1, 1)` when `layout` is `NCDHW`.
"""
def __init__(self, layout='NCDHW', **kwargs):
assert layout == 'NCDHW', "Only supports 'NCDHW' layout for now"
assert layout in ('NCDHW', 'NDHWC'),\
"Only NCDHW and NDHWC layouts are valid for 3D Pooling"
super(GlobalAvgPool3D, self).__init__(
(1, 1, 1), None, 0, True, True, 'avg', **kwargs)
(1, 1, 1), None, 0, True, True, 'avg', layout, **kwargs)


class ReflectionPad2D(HybridBlock):
Expand Down
Loading

0 comments on commit 5adb6fc

Please sign in to comment.