| |
| |
|
|
| import argparse |
| import joblib |
|
|
| from skl2onnx import convert_sklearn |
| from skl2onnx.common.data_types import FloatTensorType |
|
|
|
|
| def load_dataset(dataset_name): |
| if dataset_name == 'california': |
| from sklearn.datasets import fetch_california_housing |
| dataset = fetch_california_housing() |
| elif dataset_name == 'diabetes': |
| from sklearn.datasets import load_diabetes |
| dataset = load_diabetes() |
| elif dataset_name == 'iris': |
| from sklearn.datasets import load_iris |
| dataset = load_iris() |
| elif dataset_name == "cardiotocography": |
| from sklearn.datasets import fetch_openml |
| dataset = fetch_openml(name=dataset_name, version=1, as_frame=False) |
| X, y = dataset.data, dataset.target |
| s = y == "3" |
| y = s.astype(int) |
| return X, y |
| else: |
| raise ValueError("Invalid dataset name") |
| return dataset.data, dataset.target |
|
|
|
|
| def prepare_onnx_conversion_params(X, target_opset, model): |
| if target_opset in {9, 17}: |
| tensor = FloatTensorType([None, 4]) |
| else: |
| tensor = FloatTensorType([None, X.shape[1]]) |
| if target_opset == 9: |
| options = {id(model): {'zipmap': False}} |
| else: |
| options = None |
| return tensor, options |
|
|
|
|
| def convert2onnx(model, initial_type, options, target_opset, onnx_filename): |
| try: |
| model.save_model(onnx_filename, format="onnx") |
| print("Model saved in ONNX format successfully.") |
| except Exception as e: |
| print("Error occurred while saving model in ONNX format:", e) |
| print("Trying a second method...") |
| try: |
| onnx_model = convert_sklearn(model, initial_types=initial_type, options=options, target_opset=target_opset) |
| with open(onnx_filename, "wb") as f: |
| f.write(onnx_model.SerializeToString()) |
| print("Model converted to ONNX format and saved successfully.") |
| except Exception as e: |
| print("Error occurred while converting model to ONNX format:", e) |
|
|
| """ |
| python convert2onnx.py california adaboost_regressor.joblib adaboost_regressor.onnx |
| """ |
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description='Converts a sklearn model to ONNX format.') |
| parser.add_argument('dataset_name', type=str, help='Name of the dataset. Choose from: "california", "diabetes", "iris" or "cardiotocography".') |
| parser.add_argument('model_path', type=str, help='Path to the trained model file.') |
| parser.add_argument('onnx_filename', type=str, help='The filename for the output ONNX file.') |
| args = parser.parse_args() |
|
|
| X, y = load_dataset(args.dataset_name) |
| model = joblib.load(args.model_path) |
| target_opset = 12 |
| tensor, options = prepare_onnx_conversion_params(X, target_opset, model) |
|
|
| input_name = 'float_input' |
| initial_type = [(input_name, tensor)] |
| convert2onnx(model, initial_type, options, target_opset, args.onnx_filename) |
|
|
