def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: _size_2_t,
stride: _size_2_t = 1,
padding: _size_2_t = 0,
dilation: _size_2_t = 1,
groups: int = 1,
bias: bool = True,
padding_mode: str = "zeros", # TODO: refine this type
):
super().__init__()
assert padding_mode == "zeros"
kernel_size = _pair(kernel_size)
stride = _pair(stride)
padding = _pair(padding)
dilation = _pair(dilation)
self.groups = groups
self.weight = flow.nn.Parameter(
flow.Tensor(out_channels, in_channels // groups, *kernel_size))
self.bias = None
self._bias_add_op = None
if bias:
self.bias = flow.nn.Parameter(flow.Tensor(out_channels))
self._bias_add_op = (flow.builtin_op("bias_add").Input("a").Input(
"b").Output("out").Attr("axis", 1).Build())
self._op = (flow.builtin_op("conv2d").Input("in").Input("weight").Attr(
"filters", out_channels).Attr("padding_before", padding).Attr(
"strides", stride).Attr("kernel_size", kernel_size).Attr(
"dilation_rate", dilation).Attr("groups", groups).Attr(
"data_format", "channels_first").Output("out").Build())
self.reset_parameters()
def __init__(
self,
kernel_size: _size_2_t,
stride: Optional[_size_2_t] = None,
padding: _size_2_t = 0,
dilation: _size_2_t = 1,
return_indices: bool = False,
ceil_mode: bool = False,
):
super().__init__()
kernel_size = _pair(kernel_size)
strides = _pair(stride) if (stride is not None) else kernel_size
data_format = "NCHW"
channel_pos = "channels_last" if data_format == "NHWC" else "channels_first"
assert return_indices is False, "Only support return_indices==False for now!"
assert dilation == 1 or dilation == (1, 1), "Only support dilation==1 for now!"
padding = _pair(padding)
if len(padding) == 2:
if data_format == "NCHW":
padding = (0, 0, padding[0], padding[1])
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), 2
)
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_2d")
.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()
)