From 48c9bf5b8b0a5009d0e6c76e3c480f284ec46571 Mon Sep 17 00:00:00 2001
From: Andri Joos <andri@joos.io>
Date: Wed, 4 Oct 2023 19:28:36 +0200
Subject: [PATCH] add tf_lite_export
---
export_import/tf_lite_export.py | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
create mode 100644 export_import/tf_lite_export.py
diff --git a/export_import/tf_lite_export.py b/export_import/tf_lite_export.py
new file mode 100644
index 0000000..b8f6e0e
--- /dev/null
+++ b/export_import/tf_lite_export.py
@@ -0,0 +1,17 @@
+import tensorflow as tf
+import sys
+import os
+import tensorflow_probability as tfp # fails otherwise
+
+policies_dir = "out/policies"
+policy_name = sys.argv[1]
+policy_dir = os.path.join(policies_dir, policy_name)
+
+converter = tf.lite.TFLiteConverter.from_saved_model(policy_dir, signature_keys=["action"])
+tflite_policy = converter.convert()
+
+tflite_dir = os.path.join(policies_dir, '{}.tflite'.format(policy_name))
+with open(tflite_dir, 'wb') as f:
+ f.write(tflite_policy)
+
+print(tflite_dir)
--
GitLab