소스 검색

Merge converter and onnx script

Lutz Roeder 7 년 전
부모
커밋
5b684fd7db
4개의 변경된 파일9개의 추가작업 그리고 67개의 파일을 삭제
  1. 0 39
      tools/converter/onnx-converter
  2. 0 26
      tools/converter/onnx-converter.py
  3. 7 1
      tools/onnx
  4. 2 1
      tools/onnx-script.py

+ 0 - 39
tools/converter/onnx-converter

@@ -1,39 +0,0 @@
-#!/bin/bash
-
-set -e
-
-root=$(cd $(dirname ${0})/../..; pwd)
-build=${root}/build
-tools=${root}/tools
-third_party=${root}/third_party
-
-python=${python:-python}
-pip=${pip:-pip}
-
-identifier=onnxmltools
-repository=https://github.com/onnx/${identifier}.git
-
-mkdir -p ${third_party}
-
-if [ -d "${third_party}/${identifier}" ]; then
-    git -C "${third_party}/${identifier}" fetch -p --quiet
-    git -C "${third_party}/${identifier}" reset --quiet --hard origin/master
-else
-    echo "Clone ${repository}..."
-    git -C "${third_party}" clone --recursive ${repository}
-fi
-
-virtualenv=${build}/virtualenv/${identifier}
-virtualenv -p ${python} ${virtualenv}
-source ${virtualenv}/bin/activate
-
-pip install --quiet coremltools
-pip install --quiet tensorflow
-pip install --quiet keras
-pip install --quiet sklearn
-pip install --quiet lightgbm
-pip install --quiet ${third_party}/${identifier}
-
-python ${tools}/converter/onnx-converter.py $@
-
-deactivate

+ 0 - 26
tools/converter/onnx-converter.py

@@ -1,26 +0,0 @@
-#!/usr/bin/env python
-
-import os
-import sys
-
-file = sys.argv[1];
-base, extension = os.path.splitext(file)
-
-if extension == '.mlmodel':
-	import coremltools
-	import onnxmltools
-	coreml_model = coremltools.utils.load_spec(file)
-	onnx_model = onnxmltools.convert.convert_coreml(coreml_model)
-	onnxmltools.utils.save_model(onnx_model, base + '.onnx')
-elif extension == '.h5':
-	import keras
-	import onnxmltools
-	keras_model = keras.models.load_model(file)
-	onnx_model = onnxmltools.convert.convert_keras(keras_model)
-	onnxmltools.utils.save_model(onnx_model, base + '.onnx')
-elif extension == '.pkl':
-	from sklearn.externals import joblib
-	import onnxmltools
-	sklearn_model = joblib.load(file)
-	onnx_model = onnxmltools.convert.convert_sklearn(sklearn_model)
-	onnxmltools.utils.save_model(onnx_model, base + '.onnx')

+ 7 - 1
tools/onnx

@@ -61,8 +61,13 @@ metadata() {
 
 convert() {
     source ${virtualenv}/bin/activate
+    ${pip} install --quiet coremltools
+    ${pip} install --quiet tensorflow
+    ${pip} install --quiet keras
+    ${pip} install --quiet sklearn
+    ${pip} install --quiet lightgbm
     ${pip} install --quiet ${third_party}/onnxmltools
-
+    ${python} ${tools}/onnx-script.py convert ${1}
     deactivate
 }
 
@@ -74,5 +79,6 @@ while [ "$#" != 0 ]; do
         "build") build;;
         "schema") schema;;
         "metadata") metadata;;
+        "convert") convert ${1} && shift;;
     esac
 done

+ 2 - 1
tools/onnx-script.py

@@ -6,6 +6,7 @@ print(onnx.__file__)
 
 import json
 import io
+import os
 import sys
 
 from onnx import defs
@@ -229,6 +230,6 @@ def convert():
             text_file.write(text)
 
 if __name__ == '__main__':
-    command_table = { 'metadata': metadata }
+    command_table = { 'metadata': metadata, 'convert': convert }
     command = sys.argv[1];
     command_table[command]()