diff --git a/cli/src/dump.rs b/cli/src/dump.rs index 38eccc26ac..4756e4ec42 100644 --- 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::().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::(); let assert = diff --git a/cli/src/main.rs b/cli/src/main.rs index 4065d6e1d2..6c27ef326a 100644 --- 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")