Coverage for mlprodict/onnxrt/ops_cpu/op_average_pool.py: 92%

1# -*- encoding: utf-8 -*- 

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import itertools 

8import numpy 

9from ._op import OpRun 

10 

11 

12def _get_pad_shape(auto_pad, input_spatial_shape, kernel_spatial_shape, 

13 strides_spatial, output_spatial_shape): 

14 pad_shape = [0] * len(input_spatial_shape) 

15 if auto_pad in ('SAME_UPPER', 'SAME_LOWER'): 

16 for i in range(len(input_spatial_shape)): # pylint: disable=C0200 

17 pad_shape[i] = ( 

18 (output_spatial_shape[i] - 1) * strides_spatial[i] + 

19 kernel_spatial_shape[i] - input_spatial_shape[i]) 

20 elif auto_pad == 'VALID': 

21 pass 

22 if len(pad_shape) == 0: 

23 raise RuntimeError( # pragma: no cover 

24 "Unable to compute pad shape, auto_pad=%r, " 

25 "input_spatial_shape=%r, kernel_spatial_shape=%r, " 

26 "strides_spatial=%r." % ( 

27 auto_pad, input_spatial_shape, kernel_spatial_shape, 

28 strides_spatial)) 

29 return pad_shape 

30 

31 

32def _get_output_shape_no_ceil(auto_pad, input_spatial_shape, kernel_spatial_shape, 

33 strides_spatial): 

34 out_shape = [0] * len(input_spatial_shape) 

35 if auto_pad in ('SAME_UPPER', 'SAME_LOWER'): 

36 for i in range(len(input_spatial_shape)): # pylint: disable=C0200 

37 out_shape[i] = int( 

38 numpy.ceil( 

39 float(input_spatial_shape[i]) / 

40 float(strides_spatial[i]))) 

41 elif auto_pad == 'VALID': 

42 for i in range(len(input_spatial_shape)): # pylint: disable=C0200 

43 out_shape[i] = int( 

44 numpy.ceil( 

45 float(input_spatial_shape[i] - 

46 (kernel_spatial_shape[i] - 1)) / 

47 float(strides_spatial[i]))) 

48 return out_shape 

49 

50 

51def _get_output_shape(auto_pad, input_spatial_shape, kernel_spatial_shape, 

52 strides_spatial, pad_shape=None, ceil_mode=0): 

53 if not ceil_mode: 

54 out_shape = _get_output_shape_no_ceil( 

55 auto_pad, input_spatial_shape, kernel_spatial_shape, 

56 strides_spatial) 

57 else: 

58 round_fct = numpy.ceil if ceil_mode else numpy.floor 

59 out_shape = [0] * len(input_spatial_shape) 

60 if auto_pad in ('SAME_UPPER', 'SAME_LOWER'): 

61 for i in range(len(input_spatial_shape)): # pylint: disable=C0200 

62 out_shape[i] = int( 

63 round_fct(float(input_spatial_shape[i]) / float(strides_spatial[i]))) 

64 elif auto_pad == 'VALID': 

65 if pad_shape is None: 

66 raise ValueError( # pragma: no cover 

67 "pad_shape cannot be None if auto_pad is " 

68 "'VALID' and ceil_mode is 1.") 

69 for i in range(len(input_spatial_shape)): # pylint: disable=C0200 

70 out_shape[i] = int( 

71 round_fct( 

72 float(input_spatial_shape[i] + pad_shape[i] - kernel_spatial_shape[i]) / 

73 float(strides_spatial[i]) + 1)) 

74 if len(out_shape) == 0: 

75 raise RuntimeError( # pragma: no cover 

76 "Unable to compute output shape, auto_pad=%r, " 

77 "input_spatial_shape=%r, kernel_spatial_shape=%r, " 

78 "strides_spatial=%r, ceil_mode=%r." % ( 

79 auto_pad, input_spatial_shape, kernel_spatial_shape, 

80 strides_spatial, ceil_mode)) 

81 if min(out_shape) <= 0: 

82 raise RuntimeError( # pragma: no cover 

83 "output shape cannot be null or negative, out_shape=%r, " 

84 "auto_pad=%r, input_spatial_shape=%r, " 

85 "kernel_spatial_shape=%r, strides_spatial=%r, ceil_mode=%r." % ( 

86 out_shape, auto_pad, input_spatial_shape, 

87 kernel_spatial_shape, strides_spatial, ceil_mode)) 

88 return out_shape 

89 

90 

91def _pool(padded, x_shape, kernel_shape, strides_shape, 

92 out_shape, pad_shape, pooling_type, count_include_pad=0, ceil_mode=0): 

93 if pooling_type == 'AVG': 

94 fpool = numpy.average 

95 elif pooling_type == 'MAX': 

96 fpool = numpy.max 

97 else: 

98 raise NotImplementedError( # pragma: no cover 

99 f'Pooling type {pooling_type} does not support. Should be AVG, MAX.') 

100 spatial_size = len(x_shape) - 2 

101 y = numpy.zeros([x_shape[0], x_shape[1]] + list(out_shape)) 

102 round_fct = numpy.ceil if ceil_mode else numpy.floor 

103 

104 def loop_range(): 

105 return [range(int(round_fct( 

106 float(x_shape[i + 2] + pad_shape[i] - kernel_shape[i]) / 

107 float(strides_shape[i]) + 1))) for i in range(spatial_size)] 

108 

109 for shape in itertools.product(range(x_shape[0]), range(x_shape[1]), *loop_range()): 

110 window = padded[shape[0], shape[1]] 

111 listi = [range(strides_shape[i] * shape[i + 2], 

112 strides_shape[i] * shape[i + 2] + kernel_shape[i]) 

113 for i in range(spatial_size)] 

114 listi2 = list(itertools.product(*listi)) 

115 values = [] 

116 for i in listi2: 

117 try: 

118 values.append(window[i]) 

119 except IndexError: 

120 continue 

121 window_vals = numpy.array(values) 

122 

123 if count_include_pad == 1 and pooling_type == 'AVG': 

124 y[shape] = fpool(window_vals) 

125 else: 

126 y[shape] = fpool( 

127 window_vals[numpy.where(~numpy.isnan(window_vals))]) 

128 return y.astype(numpy.float32) 

129 

130 

131class AveragePool(OpRun): 

132 

133 atts = {'auto_pad': b'NOTSET', 

134 'ceil_mode': 0, 

135 'count_include_pad': 0, 

136 'kernel_shape': [], 

137 'pads': [], 

138 'strides': []} 

139 

140 def __init__(self, onnx_node, desc=None, **options): 

141 OpRun.__init__(self, onnx_node, desc=desc, 

142 expected_attributes=AveragePool.atts, 

143 **options) 

144 

145 def _run(self, x, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

146 if len(self.strides) == 0: 

147 strides = [1] * (len(x.shape) - 2) 

148 else: 

149 strides = self.strides 

150 kernel_shape = list(self.kernel_shape) 

151 auto_pad = ( 

152 'VALID' if self.auto_pad == b'NOTSET' 

153 else self.auto_pad.decode('ascii')) 

154 

155 if len(self.pads) == 0: 

156 pad_shape = [0] * (len(x.shape) - 2) 

157 x_shape = x.shape[2:] 

158 padded = x 

159 elif len(self.pads) == 4: 

160 pad_top, pad_bottom, pad_left, pad_right = self.pads 

161 pad_shape = [pad_top + pad_bottom, pad_left + pad_right] 

162 x_shape = numpy.array(x.shape[2:]) + numpy.array(pad_shape) 

163 const = numpy.nan if self.count_include_pad == 0 else 0 

164 padded = numpy.pad( 

165 x, ((0, 0), (0, 0), 

166 (pad_top, pad_bottom), (pad_left, pad_right)), 

167 mode='constant', constant_values=const) 

168 else: 

169 pad_shape = self.pads 

170 x_shape = x.shape[2:] 

171 padded = x 

172 

173 if auto_pad in ('SAME_LOWER', 'SAME_UPPER'): 

174 const = numpy.nan if self.count_include_pad == 0 else 0 

175 out_shape = _get_output_shape( 

176 auto_pad, x_shape, kernel_shape, strides, pad_shape, self.ceil_mode) 

177 pad_shape = _get_pad_shape( 

178 auto_pad, x_shape, kernel_shape, strides, out_shape) 

179 if auto_pad == 'SAME_LOWER': 

180 pad_bottom = pad_shape[0] // 2 

181 pad_top = pad_shape[0] - pad_bottom 

182 pad_right = pad_shape[1] // 2 

183 pad_left = pad_shape[1] - pad_right 

184 else: 

185 pad_top = pad_shape[0] // 2 

186 pad_bottom = pad_shape[0] - pad_top 

187 pad_left = pad_shape[1] // 2 

188 pad_right = pad_shape[1] - pad_left 

189 padded = numpy.pad( 

190 padded, ((0, 0), (0, 0), (pad_top, pad_bottom), 

191 (pad_left, pad_right)), 

192 mode='constant', constant_values=const) 

193 else: 

194 out_shape = _get_output_shape( 

195 auto_pad, x_shape, kernel_shape, strides, pad_shape, self.ceil_mode) 

196 

197 pooling_type = 'AVG' 

198 res = _pool(padded, x.shape, kernel_shape, strides, 

199 out_shape, pad_shape, pooling_type, 

200 count_include_pad=self.count_include_pad, 

201 ceil_mode=self.ceil_mode) 

202 return (res, )