added model conversion script.

code2k13 · Nov 4, 2022 · f30121b · f30121b
1 parent d3fbf20
commit f30121b
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+model_js
+node_modules
diff --git a/convert_model.py b/convert_model.py
@@ -0,0 +1,19 @@
+import tensorflow as tf
+import model
+import tensorflowjs as tfjs
+
+G2 = model.Generator()
+G2.trainable = False
+G2.load_weights("weights/weights")
+
+# Convert the model
+converter = tf.lite.TFLiteConverter.from_keras_model(G2) # path to the SavedModel directory
+#converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_LATENCY]
+converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
+tflite_model = converter.convert()
+
+# Save the model.
+with open('model.tflite', 'wb') as f:
+  f.write(tflite_model)
+
+tfjs.converters.save_keras_model(G2, "model_js")