saving to tflite from cli

sonos · Sep 7, 2023 · bffb4b8 · bffb4b8
1 parent 67f58a3
commit bffb4b8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/cli/src/dump.rs b/cli/src/dump.rs
@@ -220,6 +220,18 @@ pub fn handle(
         }
     }
 
+    if let Some(path) = sub_matches.value_of("tflite") {
+        let tflite = tract_tflite::tflite();
+        if let Some(mut typed) = model.downcast_ref::<TypedModel>().cloned() {
+            rename_outputs(&mut typed, sub_matches)?;
+            let file = std::fs::File::create(path)?;
+            tflite.write(&typed, file)
+                .context("Writting model to tflite")?;
+        } else {
+            bail!("Only typed model can be dumped")
+        }
+    }
+
     if options.cost {
         let total = annotations.tags.values().sum::<NodeTags>();
         let assert =

diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -328,6 +328,12 @@ fn dump_subcommand<'a>() -> clap::Command<'a> {
             .long("nnef")
             .help("Dump the network in NNEF format (as a tar.gz file)"),
             )
+        .arg(
+            Arg::new("tflite")
+            .takes_value(true)
+            .long("tflite")
+            .help("Dump the network in TfLite format"),
+            )
         .arg(
             Arg::new("compress-submodels")
             .long("compress-submodels")