MCPcopy Create free account
hub / github.com/apache/tvm / _conv2d_nchw_python

Function _conv2d_nchw_python

python/tvm/topi/testing/conv2d_nchw_python.py:27–77  ·  view source on GitHub ↗

Convolution operator in NCHW layout. Parameters ---------- a_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] w_np : numpy.ndarray 4-D with shape [num_filter, in_channel, filter_height, filter_width] stride : int or a list/tuple of two

(a_np, w_np, stride, padding)

Source from the content-addressed store, hash-verified

25
26
27def _conv2d_nchw_python(a_np, w_np, stride, padding):
28 """Convolution operator in NCHW layout.
29
30 Parameters
31 ----------
32 a_np : numpy.ndarray
33 4-D with shape [batch, in_channel, in_height, in_width]
34
35 w_np : numpy.ndarray
36 4-D with shape [num_filter, in_channel, filter_height, filter_width]
37
38 stride : int or a list/tuple of two ints
39 Stride size, or [stride_height, stride_width]
40
41 padding : int or str or a list/tuple of 2 or 4 ints
42 Padding size, or ['VALID', 'SAME'], or
43 [pad_height, pad_width] for 2 ints, or
44 [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
45
46 Returns
47 -------
48 b_np : np.ndarray
49 4-D with shape [batch, out_channel, out_height, out_width]
50 """
51 batch, in_channel, in_height, in_width = a_np.shape
52 num_filter, _, kernel_h, kernel_w = w_np.shape
53 if isinstance(stride, int):
54 stride_h = stride_w = stride
55 else:
56 stride_h, stride_w = stride
57 pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
58 pad_h = pad_top + pad_bottom
59 pad_w = pad_left + pad_right
60 # compute the output shape
61 out_channel = num_filter
62 out_height = (in_height - kernel_h + pad_h) // stride_h + 1
63 out_width = (in_width - kernel_w + pad_w) // stride_w + 1
64 b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=a_np.dtype)
65 # computation
66 for n in range(batch):
67 for f in range(out_channel):
68 for c in range(in_channel):
69 if pad_h > 0 or pad_w > 0:
70 apad = np.zeros((in_height + pad_h, in_width + pad_w), dtype=a_np.dtype)
71 apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = a_np[n, c]
72 else:
73 apad = a_np[n, c]
74
75 out = _conv2d_hw(apad, w_np[f, c])
76 b_np[n, f] += out[::stride_h, ::stride_w]
77 return b_np
78
79
80def _conv2d_hw(apad, w_np_fc):

Callers 1

conv2d_nchw_pythonFunction · 0.85

Calls 3

get_pad_tupleFunction · 0.90
_conv2d_hwFunction · 0.85
zerosMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…