export api: save tflite file

uTensor · Aug 27, 2020 · 671f2c6 · 671f2c6
1 parent 3c11a56
commit 671f2c6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/utensor_cgen/api/export.py b/utensor_cgen/api/export.py
@@ -13,6 +13,7 @@ def tflm_keras_export(
   optimizations=None,
   config_file='utensor_cli.toml',
   target='utensor',
+  output_tflite_fname=None
 ):
   with tempfile.TemporaryDirectory(prefix='utensor_') as tmp_dir:
     dir_path = Path(tmp_dir)
@@ -38,3 +39,6 @@ def tflm_keras_export(
       fid.write(tflm_model_content)
       fid.flush()
       convert_graph(fid.name, config=config_file, model_name=model_name, target=target)
+  if output_tflite_fname:
+    with open(output_tflite_fname, 'wb') as fid:
+      fid.write(tflm_model_content)