Coverage for mlprodict/onnxrt/ops_cpu/op_global_average_pool.py: 100%

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

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import numpy 

8from ..shape_object import ShapeObject 

9from ._op import OpRun 

10 

11 

12def _global_average_pool(x): 

13 spatial_shape = numpy.ndim(x) - 2 

14 y = numpy.average( 

15 x, axis=tuple(range(spatial_shape, spatial_shape + 2))) 

16 for _ in range(spatial_shape): 

17 y = numpy.expand_dims(y, -1) 

18 return y 

19 

20 

21class GlobalAveragePool(OpRun): 

22 

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

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

25 **options) 

26 

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

28 res = _global_average_pool(x) 

29 return (res, ) 

30 

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

32 if x.shape is None: 

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

34 shape = x.shape[:2] + (1, ) * (len(x.shape) - 2) 

35 return (ShapeObject(shape, dtype=x.dtype), ) 

36 

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

38 return (x, ) 

39 

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

41 res = self.run(*args) 

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