-
Notifications
You must be signed in to change notification settings - Fork 178
/
test_cml_OneHotEncoderConverter.py
56 lines (47 loc) · 2.22 KB
/
test_cml_OneHotEncoderConverter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
"""
Main functions to convert machine learned model from *Core ML* model to *ONNX*.
"""
import os
import unittest
import warnings
import numpy
try:
from sklearn.impute import SimpleImputer as Imputer
import sklearn.preprocessing
if not hasattr(sklearn.preprocessing, 'Imputer'):
# coremltools 3.1 does not work with scikit-learn 0.22
setattr(sklearn.preprocessing, 'Imputer', Imputer)
except ImportError:
from sklearn.preprocessing import Imputer
import coremltools
from sklearn.preprocessing import OneHotEncoder
from onnxmltools.convert.coreml.convert import convert
from onnxmltools.utils.tests_helper import dump_data_and_model
class TestCoremlOneHotEncoderConverter(unittest.TestCase):
def test_one_hot_encoder(self):
script_dir = os.path.dirname(__file__)
relative_path = "../data/onehot_simple.mlmodel"
abs_file = os.path.join(script_dir, relative_path)
model_coreml = coremltools.utils.load_spec(abs_file)
model_onnx = convert(model_coreml)
self.assertTrue(model_onnx is not None)
@unittest.skip('broken with the dump_data_and_model change.')
def test_conversion_one_column(self):
scikit_data = [[0], [1], [2], [4], [3], [2], [4], [5], [6], [7]]
scikit_data = numpy.asarray(scikit_data, dtype='d')
scikit_data_multiple_cols = [[0, 1], [1, 0], [2, 2], [3, 3], [4, 4]]
scikit_data_multiple_cols = numpy.asarray(scikit_data_multiple_cols, dtype='d')
scikit_model = OneHotEncoder()
# scikit_model.fit(scikit_data)
# model_coreml = coremltools.converters.sklearn.convert(scikit_model, 'single_feature', 'out')
scikit_model.fit(scikit_data_multiple_cols)
try:
model_coreml = coremltools.converters.sklearn.convert(scikit_model, ['feature_1', 'feature_2'], 'out')
except Exception as e:
warnings.warn("Unable to run convert OneHotEncoder with coreml.")
return
model_onnx = convert(model_coreml)
self.assertTrue(model_onnx is not None)
dump_data_and_model(scikit_data, scikit_model, model_onnx, basename="CmlOneHotEncoder-SkipDim1")
if __name__ == "__main__":
unittest.main()