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

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

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import numpy 

8from ._op import OpRun 

9 

10 

11class DepthToSpace(OpRun): 

12 

13 atts = {'blocksize': 0, 'mode': b'DCR'} 

14 

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

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

17 expected_attributes=DepthToSpace.atts, 

18 **options) 

19 

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

21 if len(data.shape) != 4: 

22 raise RuntimeError( # pragma: no cover 

23 f"Unexpected shape {data.shape!r}.") 

24 b, c, h, w = data.shape 

25 if self.mode == b'DCR': 

26 tmpshape = (b, self.blocksize, self.blocksize, 

27 c // (self.blocksize * self.blocksize), h, w) 

28 reshaped = data.reshape(tmpshape) 

29 transposed = numpy.transpose(reshaped, [0, 3, 4, 1, 5, 2]) 

30 else: 

31 # assert mode == "CRD" 

32 tmpshape = (b, c // (self.blocksize * self.blocksize), 

33 self.blocksize, self.blocksize, h, w) 

34 reshaped = data.reshape(tmpshape) 

35 transposed = numpy.transpose(reshaped, [0, 1, 4, 2, 5, 3]) 

36 finalshape = (b, c // (self.blocksize * self.blocksize), 

37 h * self.blocksize, w * self.blocksize) 

38 y = numpy.reshape(transposed, finalshape) 

39 return (y, ) 

40 

41 

42class SpaceToDepth(OpRun): 

43 

44 atts = {'blocksize': 0} 

45 

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

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

48 expected_attributes=SpaceToDepth.atts, 

49 **options) 

50 

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

52 if len(data.shape) != 4: 

53 raise RuntimeError( # pragma: no cover 

54 f"Unexpected shape {data.shape!r}.") 

55 b, C, H, W = data.shape 

56 tmpshape = (b, C, H // self.blocksize, self.blocksize, 

57 W // self.blocksize, self.blocksize) 

58 reshaped = numpy.reshape(data, tmpshape) 

59 transposed = numpy.transpose(reshaped, [0, 3, 5, 1, 2, 4]) 

60 finalshape = (b, C * self.blocksize * self.blocksize, 

61 H // self.blocksize, W // self.blocksize) 

62 y = numpy.reshape(transposed, finalshape) 

63 return (y, )