Coverage for mlprodict/onnx_conv/sklconv/tree_converters.py: 95%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

124 statements  

1""" 

2@file 

3@brief Rewrites some of the converters implemented in 

4:epkg:`sklearn-onnx`. 

5""" 

6import logging 

7import numpy 

8from onnx import TensorProto 

9from onnx.helper import make_attribute 

10from onnx.numpy_helper import from_array, to_array 

11from onnx.defs import onnx_opset_version 

12from skl2onnx.operator_converters.decision_tree import ( 

13 convert_sklearn_decision_tree_regressor, 

14 convert_sklearn_decision_tree_classifier) 

15from skl2onnx.operator_converters.gradient_boosting import ( 

16 convert_sklearn_gradient_boosting_regressor, 

17 convert_sklearn_gradient_boosting_classifier) 

18from skl2onnx.operator_converters.random_forest import ( 

19 convert_sklearn_random_forest_classifier, 

20 convert_sklearn_random_forest_regressor_converter) 

21from skl2onnx.common.data_types import ( 

22 guess_numpy_type, FloatTensorType, DoubleTensorType) 

23 

24 

25logger = logging.getLogger('mlprodict.onnx_conv') 

26 

27 

28def _op_type_domain_regressor(dtype, opsetml): 

29 """ 

30 Defines *op_type* and *op_domain* based on `dtype`. 

31 """ 

32 if opsetml is None: 

33 from ... import __max_supported_opsets__ 

34 if onnx_opset_version() >= 16: 

35 opsetml = min(3, __max_supported_opsets__['ai.onnx.ml']) 

36 else: 

37 opsetml = min(1, __max_supported_opsets__['ai.onnx.ml']) 

38 if opsetml >= 3: 

39 return 'TreeEnsembleRegressor', 'ai.onnx.ml', 3 

40 if dtype == numpy.float32: 

41 return 'TreeEnsembleRegressor', 'ai.onnx.ml', 1 

42 if dtype == numpy.float64: 

43 return 'TreeEnsembleRegressorDouble', 'mlprodict', 1 

44 raise RuntimeError( # pragma: no cover 

45 "Unsupported dtype {}.".format(dtype)) 

46 

47 

48def _op_type_domain_classifier(dtype, opsetml): 

49 """ 

50 Defines *op_type* and *op_domain* based on `dtype`. 

51 """ 

52 if opsetml >= 3: 

53 return 'TreeEnsembleClassifier', 'ai.onnx.ml', 3 

54 if dtype == numpy.float32: 

55 return 'TreeEnsembleClassifier', 'ai.onnx.ml', 1 

56 if dtype == numpy.float64: 

57 return 'TreeEnsembleClassifierDouble', 'mlprodict', 1 

58 raise RuntimeError( # pragma: no cover 

59 "Unsupported dtype {}.".format(dtype)) 

60 

61 

62def _fix_tree_ensemble_node(scope, container, opsetml, node, dtype): 

63 """ 

64 Fixes a node for old versionsof skl2onnx. 

65 """ 

66 atts = {'base_values': 'base_values_as_tensor', 

67 'nodes_hitrates': 'nodes_hitrates_as_tensor', 

68 'nodes_values': 'nodes_values_as_tensor', 

69 'target_weights': 'target_weights_as_tensor', 

70 'class_weights': 'class_weights_as_tensor'} 

71 logger.debug('postprocess %r name=%r opsetml=%r dtype=%r', 

72 node.op_type, node.name, opsetml, dtype) 

73 if dtype == numpy.float64: 

74 # Inserting a cast operator. 

75 index = 0 if node.op_type == 'TreeEnsembleRegressor' else 1 

76 new_name = scope.get_unique_variable_name('tree_ensemble_cast') 

77 old_name = node.output[index] 

78 node.output[index] = new_name 

79 container.add_node( 

80 'Cast', [new_name], [old_name], to=TensorProto.DOUBLE, # pylint: disable=E1101 

81 name=scope.get_unique_operator_name('tree_ensemble_cast')) 

82 attributes = list(node.attribute) 

83 del node.attribute[:] 

84 for att in attributes: 

85 if att.name in atts: 

86 logger.debug('+ rewrite att %r into %r', att.name, atts[att.name]) 

87 if att.type == 6: 

88 value = from_array( 

89 numpy.array(att.floats, dtype=dtype), atts[att.name]) 

90 elif att.type == 4: 

91 value = from_array( 

92 numpy.array(att.t.double_data, dtype=dtype), atts[att.name]) 

93 else: 

94 raise NotImplementedError( 

95 "Unable to postprocess attribute name=%r type=%r " 

96 "opsetml=%r op_type=%r (value=%r)." % ( 

97 att.name, att.type, opsetml, node.op_type, att)) 

98 if to_array(value).shape[0] == 0: 

99 raise RuntimeError( 

100 "Null value from attribute (dtype=%r): %r." % (dtype, att)) 

101 node.attribute.append(make_attribute(atts[att.name], value)) 

102 else: 

103 node.attribute.append(att) 

104 

105 

106def _fix_tree_ensemble(scope, container, opsetml, dtype): 

107 if opsetml is None: 

108 from ... import __max_supported_opsets__ 

109 if onnx_opset_version() >= 16: 

110 opsetml = min(3, __max_supported_opsets__['ai.onnx.ml']) 

111 else: 

112 opsetml = min(1, __max_supported_opsets__['ai.onnx.ml']) 

113 if opsetml < 3 or dtype == numpy.float32: 

114 return False 

115 for node in container.nodes: 

116 if node.op_type not in {'TreeEnsembleRegressor', 'TreeEnsembleClassifier'}: 

117 continue 

118 _fix_tree_ensemble_node(scope, container, opsetml, node, dtype) 

119 container.node_domain_version_pair_sets.add(('ai.onnx.ml', opsetml)) 

120 return True 

121 

122 

123def new_convert_sklearn_decision_tree_classifier(scope, operator, container): 

124 """ 

125 Rewrites the converters implemented in 

126 :epkg:`sklearn-onnx` to support an operator supporting 

127 doubles. 

128 """ 

129 dtype = guess_numpy_type(operator.inputs[0].type) 

130 if dtype != numpy.float64: 

131 dtype = numpy.float32 

132 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

133 if opsetml is None: 

134 opsetml = 3 if container.target_opset >= 16 else 1 

135 op_type, op_domain, op_version = _op_type_domain_classifier(dtype, opsetml) 

136 convert_sklearn_decision_tree_classifier( 

137 scope, operator, container, op_type=op_type, op_domain=op_domain, 

138 op_version=op_version) 

139 _fix_tree_ensemble(scope, container, opsetml, dtype) 

140 

141 

142def new_convert_sklearn_decision_tree_regressor(scope, operator, container): 

143 """ 

144 Rewrites the converters implemented in 

145 :epkg:`sklearn-onnx` to support an operator supporting 

146 doubles. 

147 """ 

148 dtype = guess_numpy_type(operator.inputs[0].type) 

149 if dtype != numpy.float64: 

150 dtype = numpy.float32 

151 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

152 op_type, op_domain, op_version = _op_type_domain_regressor(dtype, opsetml) 

153 convert_sklearn_decision_tree_regressor( 

154 scope, operator, container, op_type=op_type, op_domain=op_domain, 

155 op_version=op_version) 

156 _fix_tree_ensemble(scope, container, opsetml, dtype) 

157 

158 

159def new_convert_sklearn_gradient_boosting_classifier(scope, operator, container): 

160 """ 

161 Rewrites the converters implemented in 

162 :epkg:`sklearn-onnx` to support an operator supporting 

163 doubles. 

164 """ 

165 dtype = guess_numpy_type(operator.inputs[0].type) 

166 if dtype != numpy.float64: 

167 dtype = numpy.float32 

168 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

169 if opsetml is None: 

170 opsetml = 3 if container.target_opset >= 16 else 1 

171 op_type, op_domain, op_version = _op_type_domain_classifier(dtype, opsetml) 

172 convert_sklearn_gradient_boosting_classifier( 

173 scope, operator, container, op_type=op_type, op_domain=op_domain, 

174 op_version=op_version) 

175 _fix_tree_ensemble(scope, container, opsetml, dtype) 

176 

177 

178def new_convert_sklearn_gradient_boosting_regressor(scope, operator, container): 

179 """ 

180 Rewrites the converters implemented in 

181 :epkg:`sklearn-onnx` to support an operator supporting 

182 doubles. 

183 """ 

184 dtype = guess_numpy_type(operator.inputs[0].type) 

185 if dtype != numpy.float64: 

186 dtype = numpy.float32 

187 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

188 op_type, op_domain, op_version = _op_type_domain_regressor(dtype, opsetml) 

189 convert_sklearn_gradient_boosting_regressor( 

190 scope, operator, container, op_type=op_type, op_domain=op_domain, 

191 op_version=op_version) 

192 _fix_tree_ensemble(scope, container, opsetml, dtype) 

193 

194 

195def new_convert_sklearn_random_forest_classifier(scope, operator, container): 

196 """ 

197 Rewrites the converters implemented in 

198 :epkg:`sklearn-onnx` to support an operator supporting 

199 doubles. 

200 """ 

201 dtype = guess_numpy_type(operator.inputs[0].type) 

202 if dtype != numpy.float64: 

203 dtype = numpy.float32 

204 if (dtype == numpy.float64 and 

205 isinstance(operator.outputs[1].type, FloatTensorType)): 

206 operator.outputs[1].type = DoubleTensorType( 

207 operator.outputs[1].type.shape) 

208 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

209 if opsetml is None: 

210 opsetml = 3 if container.target_opset >= 16 else 1 

211 op_type, op_domain, op_version = _op_type_domain_classifier(dtype, opsetml) 

212 convert_sklearn_random_forest_classifier( 

213 scope, operator, container, op_type=op_type, op_domain=op_domain, 

214 op_version=op_version) 

215 _fix_tree_ensemble(scope, container, opsetml, dtype) 

216 

217 

218def new_convert_sklearn_random_forest_regressor(scope, operator, container): 

219 """ 

220 Rewrites the converters implemented in 

221 :epkg:`sklearn-onnx` to support an operator supporting 

222 doubles. 

223 """ 

224 dtype = guess_numpy_type(operator.inputs[0].type) 

225 if dtype != numpy.float64: 

226 dtype = numpy.float32 

227 opsetml = container.target_opset_all.get('ai.onnx.ml', None) 

228 if opsetml is None: 

229 opsetml = 3 if container.target_opset >= 16 else 1 

230 op_type, op_domain, op_version = _op_type_domain_regressor(dtype, opsetml) 

231 convert_sklearn_random_forest_regressor_converter( 

232 scope, operator, container, op_type=op_type, op_domain=op_domain, 

233 op_version=op_version) 

234 _fix_tree_ensemble(scope, container, opsetml, dtype)