コード例 #1
0
    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()
        )
コード例 #2
0
    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()
        )