def __init__(
self,
kernel_size: _size_3_t,
stride: Optional[_size_3_t] = None,
padding: _size_3_t = 0,
dilation: _size_3_t = 1,
return_indices: bool = False,
ceil_mode: bool = False,
):
super().__init__()
kernel_size = _triple(kernel_size)
strides = _triple(stride) if (stride is not None) else kernel_size
data_format = "NCDHW"
channel_pos = "channels_last" if data_format == "NDHWC" else "channels_first"
assert return_indices is False, "Only support return_indices==False for now!"
assert dilation == 1 or dilation == (
1,
1,
1,
), "Only support dilation==1 for now!"
padding = _triple(padding)
if len(padding) == 3:
if data_format == "NCDHW":
padding = (0, 0, padding[0], padding[1], padding[2])
else:
raise ValueError("error padding param!")
else:
raise ValueError("error padding param!")
padding_type, pads_list = calc_pool_padding(
padding, get_dhw_offset(channel_pos), 3
)
padding_before = [pad[0] for pad in pads_list]
padding_after = [pad[1] for pad in pads_list]
self._op = (
flow.builtin_op("max_pool_3d")
.Attr("data_format", channel_pos)
.Attr("pool_size", kernel_size)
.Attr("strides", strides)
.Attr("ceil_mode", ceil_mode)
.Attr("padding", padding_type)
.Attr("padding_before", padding_before)
.Attr("padding_after", padding_after)
.Input("x")
.Output("y")
.Build()
)
def __init__(
self,
kernel_size: _size_2_t,
stride: Optional[_size_2_t] = None,
padding: _size_2_t = 0,
ceil_mode: bool = False,
count_include_pad: Optional[bool] = None,
divisor_override: Optional[int] = None,
name: Optional[str] = None,
):
super().__init__()
kernel_size = _pair(kernel_size)
stride = _pair(stride) if (stride is not None) else kernel_size
assert isinstance(padding, int) or isinstance(
padding, tuple
), "padding can only int int or tuple of 2 ints."
padding = _pair(padding)
padding = [0, 0, *padding]
assert count_include_pad is None, "count_include_pad not supported yet"
assert divisor_override is None, "divisor_override not supported yet"
_channel_pos = "channels_first"
# TODO(yaochi): align with pytorch when padding is asymmetric
_padding_type, _pads_list = calc_pool_padding(
padding, get_dhw_offset(_channel_pos), 2
)
_padding_before = [pad[0] for pad in _pads_list]
_padding_after = [pad[1] for pad in _pads_list]
self._op = (
flow.builtin_op("avg_pool_2d", name)
.Attr("data_format", _channel_pos)
.Attr("pool_size", kernel_size)
.Attr("strides", stride)
.Attr("ceil_mode", ceil_mode)
.Attr("padding", _padding_type)
.Attr("padding_before", _padding_before)
.Attr("padding_after", _padding_after)
.Input("x")
.Output("y")
.Build()
)