File size: 8,208 Bytes
dc2106c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0


import numpy as np

from onnx.reference.op_run import OpRun
from onnx.reference.ops._op_common_indices import _get_indices, _is_out


def _col2im_shape_check_2d(X, output_shape, kernel_shape, dilations, pads, strides):  # type: ignore
    output_height, output_width = output_shape
    kernel_height, kernel_width = kernel_shape
    dilation_height, dilation_width = dilations
    stride_height, stride_width = strides

    ndim = len(X.shape)
    if not (
        (ndim == 2 and X.shape[0] != 0 and X.shape[1] != 0)
        or (ndim == 3 and X.shape[1] != 0 and X.shape[2] != 0)
    ):
        raise ValueError(
            "Expected 2D or 3D (batch mode) tensor for input with possibly 0 batch size and non-zero dimensions for input."
        )

    batch_dim = 0 if len(X.shape) == 3 else -1
    n_input_plane = X.shape[batch_dim + 1]

    if n_input_plane % (kernel_width * kernel_height) != 0:
        raise ValueError(
            f"Expected size of input's dimension 1 to be divisible by the "
            f"product of kernel_size, but got input.size(1)={n_input_plane} "
            f"and kernel_size={kernel_shape}."
        )

    input_length = X.shape[batch_dim + 2]
    n_blocks_height = (
        output_height + pads[0, :].sum() - dilation_height * (kernel_height - 1) - 1
    ) // stride_height + 1
    n_blocks_width = (
        output_width + pads[1, :].sum() - dilation_width * (kernel_width - 1) - 1
    ) // stride_width + 1

    if input_length != (n_blocks_height * n_blocks_width):
        raise ValueError(
            f"Given batch_dim={batch_dim}, n_input_plane={n_input_plane}, X.shape={X.shape}, "
            f"output_shape={output_shape}, kernel_shape={kernel_shape}, "
            f"dilations={dilations}, pads={pads}, strides={strides}, "
            f"expected size of input's dimension 2 to match the calculated number of ",
            f"sliding blocks {n_blocks_height} * {n_blocks_width} = {n_blocks_height * n_blocks_width}, "
            f"but got input.size(2)={input_length}.",
        )

    if not (n_blocks_height >= 1 and n_blocks_width >= 1):
        raise ValueError(
            f"Given batch_dim={batch_dim}, n_input_plane={n_input_plane}, X.shape={X.shape}, "
            f"output_shape={output_shape}, kernel_shape={kernel_shape}, "
            f"dilations={dilations}, pads={pads}, strides={strides}, "
            f"calculated shape of the array of sliding blocks as ({n_blocks_height}, {n_blocks_width}), "
            f"which is too small (non-positive)."
        )


def _col2im_naive_implementation_2d(res, image_shape, kernel_shape, dilations, pads, strides):  # type: ignore
    # source: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/im2col.h

    n_dims = len(pads) // 2
    new_pads = np.array([(pads[i], pads[i + n_dims]) for i in range(n_dims)])
    _col2im_shape_check_2d(res, image_shape, kernel_shape, dilations, new_pads, strides)

    data_col = res.ravel()
    data_im = np.zeros(image_shape, dtype=res.dtype).flatten()

    kernel_h, kernel_w = kernel_shape
    channels_col = kernel_h * kernel_w
    stride_h, stride_w = strides
    dilation_h, dilation_w = dilations
    pad_h, pad_w = new_pads[:, 0]
    height, width = image_shape
    output_height, output_width = image_shape

    height_col = (
        output_height + new_pads[0, :].sum() - (dilation_h * (kernel_h - 1) + 1)
    ) // stride_h + 1
    width_col = (
        output_width + new_pads[1, :].sum() - (dilation_w * (kernel_w - 1) + 1)
    ) // stride_w + 1

    for c_col in range(channels_col):
        w_offset = c_col % kernel_w
        h_offset = (c_col // kernel_w) % kernel_h
        c_im = c_col // (kernel_h * kernel_w)

        for h_col in range(height_col):
            h_im = h_col * stride_h - pad_h + h_offset * dilation_h
            for w_col in range(width_col):
                w_im = w_col * stride_w - pad_w + w_offset * dilation_w
                if 0 <= h_im < height and 0 <= w_im < width:
                    i_im = (c_im * height + h_im) * width + w_im
                    i_col = (c_col * height_col + h_col) * width_col + w_col
                    if 0 <= i_col < data_col.shape[0]:
                        data_im[i_im] += data_col[i_col]

    return data_im.reshape(image_shape)


def _col2im_shape_check(X, output_shape, kernel_shape, dilations, pads, strides):  # type: ignore
    n_input_plane = X.shape[0]

    kernel_size = np.prod(kernel_shape)

    if n_input_plane % kernel_size != 0:
        raise ValueError(
            f"Expected size of input's dimension 1 to be divisible by the "
            f"product of kernel_size={kernel_size}, "
            f"but got input.size(1)={n_input_plane} "
            f"and kernel_shape={kernel_shape}, X.shape={X.shape}, output_shape={output_shape}."
        )

    input_length = X.shape[1]
    n_dims = len(output_shape)
    n_blocks = []
    for i in range(n_dims):
        n_block = (
            output_shape[i]
            + pads[i, :].sum()
            - dilations[i] * (kernel_shape[i] - 1)
            - 1
        ) // strides[i] + 1
        n_blocks.append(n_block)

    block_size = np.prod(n_blocks)
    if input_length != block_size:
        raise ValueError(
            f"Given n_input_plane={n_input_plane}, X.shape={X.shape}, "
            f"output_shape={output_shape}, kernel_shape={kernel_shape}, "
            f"dilations={dilations}, pads={pads}, strides={strides}, "
            f"expected size of input's dimension 2 to match the calculated number of "
            f"sliding blocks {n_blocks} = {block_size}, "
            f"but got input.size(2)={input_length}.",
        )


def col2im_naive_implementation(data, image_shape, kernel_shape, dilations, pads, strides):  # type: ignore
    """Naive implementation for `col2im`."""
    n_dims = len(pads) // 2
    new_pads = np.array([(pads[i], pads[i + n_dims]) for i in range(n_dims)])
    _col2im_shape_check(data, image_shape, kernel_shape, dilations, new_pads, strides)

    data_col = data
    data_im = np.zeros(image_shape, dtype=data.dtype)

    dim_col = []
    for i in range(n_dims):
        col = (
            image_shape[i]
            + new_pads[i, :].sum()
            - (dilations[i] * (kernel_shape[i] - 1) + 1)
        ) // strides[i] + 1
        dim_col.append(col)

    kernel_size = np.prod(kernel_shape)
    col_size = np.prod(dim_col)
    for c_col in range(kernel_size):
        offset = _get_indices(c_col, kernel_shape)

        for col in range(col_size):
            ind_col = _get_indices(col, dim_col)
            ind_im = []
            for i in range(n_dims):
                ind = (
                    ind_col[i] * strides[i] - new_pads[i, 0] + offset[i] * dilations[i]
                )
                ind_im.append(ind)

            if not _is_out(ind_im, data_im.shape):
                data_im[tuple(ind_im)] += data_col[c_col, col]

    return data_im


class Col2Im(OpRun):
    def _run(self, data, image_shape, block_shape, dilations=None, pads=None, strides=None):  # type: ignore
        if dilations is None:
            dilations = [1 for s in image_shape]
        if pads is None:
            pads = [0 for s in image_shape] * 2
        if strides is None:
            strides = [1 for s in image_shape]

        bl = np.prod(block_shape)
        C = data.shape[1] // bl
        data = data.reshape(data.shape[:1] + (C,) + (bl,) + data.shape[2:])

        ks = tuple(block_shape)
        res = None
        for n in range(data.shape[0]):
            for c in range(data.shape[1]):
                out = col2im_naive_implementation(
                    data[n, c, ...], image_shape, ks, dilations, pads, strides
                )
                if res is None:
                    new_shape = data.shape[:2] + out.shape
                    res = np.empty(new_shape, dtype=data.dtype)
                res[n, c, ...] = out
        return (res,)  # type: ignore