feat(components): Added converter components for Tensorflow Lite mode…

…ls (#4532)
kubeflow · Sep 22, 2020 · 1802db3 · 1802db3
1 parent 06bf429
commit 1802db3
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/components/_converters/TensorflowLiteModel/from_KerasModelHdf5/component.yaml b/components/_converters/TensorflowLiteModel/from_KerasModelHdf5/component.yaml
@@ -0,0 +1,22 @@
+name: Convert Keras HDF5 model to Tensorflow Lite model
+inputs:
+- {name: Model, type: KerasModelHdf5}
+outputs:
+- {name: Model, type: TensorflowLiteModel}
+metadata:
+  annotations:
+    author: Alexey Volkov <alexey.volkov@ark-kun.com>
+implementation:
+  container:
+    image: tensorflow/tensorflow:2.3.0
+    command:
+    - sh
+    - -exc
+    - |
+      model_path="$0"
+      output_model_path="$1"
+      mkdir -p "$(dirname "$output_model_path")"
+      
+      tflite_convert --keras_model_file "$model_path" --output_file "$output_model_path"
+    - {inputPath: Model}
+    - {outputPath: Model}
diff --git a/components/_converters/TensorflowLiteModel/from_TensorflowSavedModel/component.yaml b/components/_converters/TensorflowLiteModel/from_TensorflowSavedModel/component.yaml
@@ -0,0 +1,22 @@
+name: Convert Tensorflow SavedModel to Tensorflow Lite model
+inputs:
+- {name: Model, type: TensorflowSavedModel}
+outputs:
+- {name: Model, type: TensorflowLiteModel}
+metadata:
+  annotations:
+    author: Alexey Volkov <alexey.volkov@ark-kun.com>
+implementation:
+  container:
+    image: tensorflow/tensorflow:2.3.0
+    command:
+    - sh
+    - -exc
+    - |
+      model_path="$0"
+      output_model_path="$1"
+      mkdir -p "$(dirname "$output_model_path")"
+      
+      tflite_convert --saved_model_dir "$model_path" --output_file "$output_model_path"
+    - {inputPath: Model}
+    - {outputPath: Model}