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

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

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import itertools 

8import numpy 

9from ..shape_object import ShapeObjectFct, ShapeObject 

10from ._op import OpRun 

11 

12 

13def _get_pad_shape(auto_pad, input_spatial_shape, kernel_spatial_shape, 

14 strides_spatial, output_spatial_shape): 

15 pad_shape = [0] * len(input_spatial_shape) 

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

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

18 pad_shape[i] = ( 

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

20 kernel_spatial_shape[i] - input_spatial_shape[i]) 

21 elif auto_pad == 'VALID': 

22 pass 

23 if len(pad_shape) == 0: 

24 raise RuntimeError( # pragma: no cover 

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

26 "input_spatial_shape=%r, kernel_spatial_shape=%r, " 

27 "strides_spatial=%r." % ( 

28 auto_pad, input_spatial_shape, kernel_spatial_shape, 

29 strides_spatial)) 

30 return pad_shape 

31 

32 

33def _get_output_shape_no_ceil(auto_pad, input_spatial_shape, kernel_spatial_shape, 

34 strides_spatial): 

35 out_shape = [0] * len(input_spatial_shape) 

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

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

38 out_shape[i] = int( 

39 numpy.ceil( 

40 float(input_spatial_shape[i]) / 

41 float(strides_spatial[i]))) 

42 elif auto_pad == 'VALID': 

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

44 out_shape[i] = int( 

45 numpy.ceil( 

46 float(input_spatial_shape[i] - 

47 (kernel_spatial_shape[i] - 1)) / 

48 float(strides_spatial[i]))) 

49 return out_shape 

50 

51 

52def _get_output_shape(auto_pad, input_spatial_shape, kernel_spatial_shape, 

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

54 if not ceil_mode: 

55 out_shape = _get_output_shape_no_ceil( 

56 auto_pad, input_spatial_shape, kernel_spatial_shape, 

57 strides_spatial) 

58 else: 

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

60 out_shape = [0] * len(input_spatial_shape) 

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

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

63 out_shape[i] = int( 

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

65 elif auto_pad == 'VALID': 

66 if pad_shape is None: 

67 raise ValueError( # pragma: no cover 

68 "pad_shape cannot be None if auto_pad is " 

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

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

71 out_shape[i] = int( 

72 round_fct( 

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

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

75 if len(out_shape) == 0: 

76 raise RuntimeError( # pragma: no cover 

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

78 "input_spatial_shape=%r, kernel_spatial_shape=%r, " 

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

80 auto_pad, input_spatial_shape, kernel_spatial_shape, 

81 strides_spatial, ceil_mode)) 

82 if min(out_shape) <= 0: 

83 raise RuntimeError( # pragma: no cover 

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

85 "auto_pad=%r, input_spatial_shape=%r, " 

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

87 out_shape, auto_pad, input_spatial_shape, 

88 kernel_spatial_shape, strides_spatial, ceil_mode)) 

89 return out_shape 

90 

91 

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

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

94 if pooling_type == 'AVG': 

95 fpool = numpy.average 

96 elif pooling_type == 'MAX': 

97 fpool = numpy.max 

98 else: 

99 raise NotImplementedError( # pragma: no cover 

100 'Pooling type {} does not support. Should be AVG, MAX.' 

101 ''.format(pooling_type)) 

102 spatial_size = len(x_shape) - 2 

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

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

105 

106 def loop_range(): 

107 return [range(int(round_fct( 

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

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

110 

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

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

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

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

115 for i in range(spatial_size)] 

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

117 values = [] 

118 for i in listi2: 

119 try: 

120 values.append(window[i]) 

121 except IndexError: 

122 continue 

123 window_vals = numpy.array(values) 

124 

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

126 y[shape] = fpool(window_vals) 

127 else: 

128 y[shape] = fpool( 

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

130 return y.astype(numpy.float32) 

131 

132 

133class AveragePool(OpRun): 

134 

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

136 'ceil_mode': 0, 

137 'count_include_pad': 0, 

138 'kernel_shape': [], 

139 'pads': [], 

140 'strides': []} 

141 

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

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

144 expected_attributes=AveragePool.atts, 

145 **options) 

146 

147 def _run(self, x): # pylint: disable=W0221 

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

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

150 else: 

151 strides = self.strides 

152 kernel_shape = list(self.kernel_shape) 

153 auto_pad = ( 

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

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

156 

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

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

159 x_shape = x.shape[2:] 

160 padded = x 

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

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

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

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

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

166 padded = numpy.pad( 

167 x, ((0, 0), (0, 0), 

168 (pad_top, pad_bottom), (pad_left, pad_right)), 

169 mode='constant', constant_values=const) 

170 else: 

171 pad_shape = self.pads 

172 x_shape = x.shape[2:] 

173 padded = x 

174 

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

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

177 out_shape = _get_output_shape( 

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

179 pad_shape = _get_pad_shape( 

180 auto_pad, x_shape, kernel_shape, strides, out_shape) 

181 if auto_pad == 'SAME_LOWER': 

182 pad_bottom = pad_shape[0] // 2 

183 pad_top = pad_shape[0] - pad_bottom 

184 pad_right = pad_shape[1] // 2 

185 pad_left = pad_shape[1] - pad_right 

186 else: 

187 pad_top = pad_shape[0] // 2 

188 pad_bottom = pad_shape[0] - pad_top 

189 pad_left = pad_shape[1] // 2 

190 pad_right = pad_shape[1] - pad_left 

191 padded = numpy.pad( 

192 padded, ((0, 0), (0, 0), (pad_top, pad_bottom), 

193 (pad_left, pad_right)), 

194 mode='constant', constant_values=const) 

195 else: 

196 out_shape = _get_output_shape( 

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

198 

199 pooling_type = 'AVG' 

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

201 out_shape, pad_shape, pooling_type, 

202 count_include_pad=self.count_include_pad, 

203 ceil_mode=self.ceil_mode) 

204 return (res, ) 

205 

206 def _infer_shapes(self, x): # pylint: disable=W0221 

207 kernel_shape = list(self.kernel_shape) 

208 auto_pad = 'VALID' if self.auto_pad == 'NOTSET' else self.auto_pad 

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

210 if x.shape is None: 

211 return (ShapeObject(None, dtype=x.dtype), ) 

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

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

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

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

216 

217 def compute_shape(xshape): 

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

219 strides = [1] * (len(xshape) - 2) 

220 else: 

221 strides = self.strides 

222 out_shape = _get_output_shape( 

223 auto_pad, xshape[2:], kernel_shape, strides, pad_shape, self.ceil_mode) 

224 return out_shape 

225 

226 return (ShapeObjectFct( 

227 compute_shape, x, name="AveragePool", dtype=x.dtype), ) 

228 

229 def _infer_types(self, x): # pylint: disable=W0221 

230 return (x, ) 

231 

232 def _infer_sizes(self, *args): # pylint: disable=W0221 

233 res = self.run(*args) 

234 return (dict(temp=0), ) + res