Coverage for mlprodict/onnxrt/ops_cpu/op_softmax.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 OpRunUnaryNum, OpRunBinaryNum 

9from ._new_ops import OperatorSchema 

10 

11 

12class Softmax(OpRunUnaryNum): 

13 

14 atts = {'axis': 1} 

15 

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

17 OpRunUnaryNum.__init__(self, onnx_node, desc=desc, 

18 expected_attributes=Softmax.atts, 

19 **options) 

20 

21 def _run(self, X): # pylint: disable=W0221 

22 if self.inplaces.get(0, False): 

23 return self._run_inplace(X) 

24 tmp = X - X.max(axis=self.axis, keepdims=1) 

25 Y = numpy.exp(tmp) 

26 Y /= Y.sum(axis=self.axis, keepdims=1) 

27 return (Y, ) 

28 

29 def _run_inplace(self, X): 

30 X -= X.max(axis=self.axis, keepdims=1) 

31 numpy.exp(X, out=X) 

32 X /= X.sum(axis=self.axis, keepdims=1) 

33 return (X, ) 

34 

35 def to_python(self, inputs): 

36 lines = ["tmp = {0} - {0}.max(axis=axis)[:, numpy.newaxis]".format( 

37 inputs[0]), 

38 "Y = numpy.exp(tmp)", 

39 "Y /= Y.sum(axis=axis)[:, numpy.newaxis]", 

40 "return Y"] 

41 return ("import numpy", "\n".join(lines)) 

42 

43 

44class SoftmaxGrad_13(OpRunBinaryNum): 

45 """ 

46 SoftmaxGrad computes :math:`dX = Y * ( dY - ReduceSum(Y * dY))`. 

47 ONNX does not have a dot product, 

48 which can be simulated as a pointwise-multiplication ("Mul"), 

49 followed by a "ReduceSum". Unfortunately, the treatment of "axis" 

50 is different in "SoftmaxGrad" and "ReduceSum". 

51 If axis=k for SoftmaxGrad, we need to specify [k, ..., n-1] as the axes of 

52 reduction for "ReduceSum", after accounting for negative-axis specification. 

53 An alternative solution would be to Flatten inputs to 2D and then reshape 

54 output back to original shape. Hopefully, many of these ops can be optimized 

55 away in the common-case of statically-known shapes. 

56 """ 

57 

58 atts = {'axis': 1} 

59 

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

61 OpRunBinaryNum.__init__(self, onnx_node, desc=desc, 

62 expected_attributes=SoftmaxGrad_13.atts, 

63 **options) 

64 

65 def _find_custom_operator_schema(self, op_name): 

66 if op_name in ("SoftmaxGrad_13", "SoftmaxGrad"): 

67 return SoftmaxGradSchema() 

68 raise RuntimeError( # pragma: no cover 

69 "Unable to find a schema for operator '{}'.".format(op_name)) 

70 

71 def _run(self, grad, prob): # pylint: disable=W0221 

72 # softmax 

73 # tmp = X - X.max(axis=self.axis)[:, numpy.newaxis] 

74 # Y = numpy.exp(tmp) 

75 # Y /= Y.sum(axis=self.axis)[:, numpy.newaxis] 

76 # derivative 

77 pg = prob * grad 

78 if self.axis < 0: 

79 axis = len(pg.shape) + self.axis 

80 else: 

81 axis = self.axis 

82 axis = tuple(range(axis, len(pg.shape))) 

83 dg = grad - pg.sum(axis=axis, keepdims=1) 

84 return (prob * dg, ) 

85 

86 

87class SoftmaxGradSchema(OperatorSchema): 

88 """ 

89 Defines a schema for operators added in this package 

90 such as @see cl SoftmaxGrad_13. 

91 """ 

92 

93 def __init__(self): 

94 OperatorSchema.__init__(self, 'SoftmaxGrad') 

95 self.attributes = SoftmaxGrad_13.atts 

96 

97 

98SoftmaxGrad = SoftmaxGrad_13