Make exporting tflite optional and add instructions to README.

Related issues: #341 #158 #138 #4

Author: mingxingtan

efficientdet/README.md

@@ -65,16 +65,36 @@ We have provided a list of EfficientDet checkpoints and results as follows:
 ** <em>val</em> denotes validation results, <em>test-dev</em> denotes test-dev2017 results. AP<sup>val</sup> is for validation accuracy, all other AP results in the table are for COCO test-dev2017. All accuracy numbers are for single-model single-scale without ensemble or test-time augmentation. All checkpoints are trained with baseline preprocessing (no autoaugmentation).
 ** EfficientNet-D0 to D6 are trained with 300 epochs, EfficientNet-D7 is trained with 500 epochs.
 
-## 3. Benchmark model latency.
+
+## 3. Export SavedModel, frozen graph, or tflite.
+
+Run the following command line to export models:
+
+    !rm  -rf savedmodeldir
+    !python model_inspect.py --runmode=saved_model --model_name=efficientdet-d0 \
+      --ckpt_path=efficientdet-d0 --saved_model_dir=savedmodeldir \
+      --tflite_path=efficientdet-d0.tflite
+
+Then you will get:
+
+ - saved model under savedmodeldir/
+ - frozen graph with name savedmodeldir/efficientdet-d0_frozen.pb
+ - tflite file with name efficientdet-d0.tflite
+
+Notably, --tflite_path is optinal, and it only works after 2.2.0-rc4.
+
+
+## 4. Benchmark model latency.
 
 
 There are two types of latency: network latency and end-to-end latency.
 
 (1) To measure the network latency (from the fist conv to the last class/box
 prediction output), use the following command:
 
-    !python model_inspect.py --runmode=bm --model_name=efficientdet-d0 \
-      # --hparams="precision=mixed-float16"  # uncomment if on V100
+    !python model_inspect.py --runmode=bm --model_name=efficientdet-d0
+
+** add --hparams="precision=mixed-float16" if running on V100.
 
 On single Tesla V100 without TensorRT, our D0 network (no pre/post-processing)
 has 134 FPS (frame per second) for batch size 1, and 238 FPS for batch size 8.
@@ -86,13 +106,11 @@ use the following command:
     !rm  -rf /tmp/benchmark/
     !python model_inspect.py --runmode=saved_model --model_name=efficientdet-d0 \
       --ckpt_path=efficientdet-d0 --saved_model_dir=/tmp/benchmark/ \
-      # --hparams="precision=mixed-float16"  # uncomment if on V100
 
     !python model_inspect.py --runmode=saved_model_benchmark \
       --saved_model_dir=/tmp/benchmark/efficientdet-d0_frozen.pb \
       --model_name=efficientdet-d0  --input_image=testdata/img1.jpg  \
       --output_image_dir=/tmp/  \
-      # --hparams="precision=mixed-float16"  # uncomment if on V100
 
 On single Tesla V100 without using TensorRT, our end-to-end
 latency and throughput are:
@@ -109,7 +127,7 @@ latency and throughput are:
 
 ** FPS means frames per second (or images/second).
 
-## 4. Inference for images.
+## 5. Inference for images.
 
     # Step0: download model and testing image.
     !export MODEL=efficientdet-d0
@@ -157,7 +175,7 @@ Here is an example of EfficientDet-D0 visualization: more on [tutorial](tutorial
 <img src="./g3doc/street.jpg" width="800" />
 </p>
 
-## 5. Inference for videos.
+## 6. Inference for videos.
 
 You can run inference for a video and show the results online:
 
@@ -180,7 +198,7 @@ You can run inference for a video and show the results online:
       --saved_model_dir=/tmp/savedmodel --input_video=input.mov  \
       --output_video=output.mov
 
-## 6. Eval on COCO 2017 val or test-dev.
+## 7. Eval on COCO 2017 val or test-dev.
 
     // Download coco data.
     !wget http://images.cocodataset.org/zips/val2017.zip
@@ -227,7 +245,7 @@ You can also run eval on test-dev set with the following command:
     # Now you can submit testdev_output/detections_test-dev2017_test_results.json to
     # coco server: https://competitions.codalab.org/competitions/20794#participate
 
-## 7. Train on PASCAL VOC 2012 with backbone ImageNet ckpt.
+## 8. Train on PASCAL VOC 2012 with backbone ImageNet ckpt.
 
     # Download and convert pascal data.
     !wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
@@ -253,7 +271,7 @@ You can also run eval on test-dev set with the following command:
         --hparams="num_classes=20,moving_average_decay=0" \
         --use_tpu=False
 
-## 8. Finetune on PASCAL VOC 2012 with detector COCO ckpt.
+## 9. Finetune on PASCAL VOC 2012 with detector COCO ckpt.
 Create a config file for the PASCAL VOC dataset called voc_config.yaml and put this in it.
 
       num_classes: 20
@@ -289,7 +307,7 @@ If you want to do inference for custom data, you can run
 
 You should check more details of runmode which is written in caption-4.
 
-## 9. Training EfficientDets on TPUs.
+## 10. Training EfficientDets on TPUs.
 
 To train this model on Cloud TPU, you will need:
 

efficientdet/inference.py

@@ -699,7 +699,7 @@ def to_tflite(self, saved_model_dir):
     converter.target_spec.supported_ops = supported_ops
     return converter.convert()
 
-  def export(self, output_dir, frozen_pb=True, tflite=True):
+  def export(self, output_dir, frozen_pb=True, tflite_path=None):
     """Export a saved model."""
     signitures = self.signitures
     signature_def_map = {
@@ -729,12 +729,11 @@ def export(self, output_dir, frozen_pb=True, tflite=True):
       tf.io.gfile.GFile(pb_path, 'wb').write(graphdef.SerializeToString())
       logging.info('Free graph saved at %s', pb_path)
 
-    if tflite:
+    if tflite_path:
       ver = tf.__version__
       if ver < '2.2.0-dev20200501' or ('dev' not in ver and ver < '2.2.0-rc4'):
         raise ValueError('TFLite requires TF 2.2.0rc4 or laterr version.')
       tflite_model = self.to_tflite(output_dir)
-      tflite_path = os.path.join(output_dir, self.model_name + '.tflite')
       with tf.io.gfile.GFile(tflite_path, 'wb') as f:
         f.write(tflite_model)
 

efficientdet/model_inspect.py

@@ -71,6 +71,7 @@
 # For saved model.
 flags.DEFINE_string('saved_model_dir', '/tmp/saved_model',
                     'Folder path for saved model.')
+flags.DEFINE_string('tflite_path', None, 'Path for exporting tflite file.')
 
 FLAGS = flags.FLAGS
 
@@ -86,6 +87,7 @@ def __init__(self,
                ckpt_path: Text = None,
                export_ckpt: Text = None,
                saved_model_dir: Text = None,
+               tflite_path: Text = None,
                batch_size: int = 1,
                hparams: Text = ''):
     self.model_name = model_name
@@ -95,6 +97,7 @@ def __init__(self,
     self.ckpt_path = ckpt_path
     self.export_ckpt = export_ckpt
     self.saved_model_dir = saved_model_dir
+    self.tflite_path = tflite_path
 
     model_config = hparams_config.get_detection_config(model_name)
     model_config.override(hparams)  # Add custom overrides
@@ -144,7 +147,7 @@ def export_saved_model(self, **kwargs):
         model_params=self.model_config.as_dict(),
         **kwargs)
     driver.build()
-    driver.export(self.saved_model_dir)
+    driver.export(self.saved_model_dir, tflite_path=self.tflite_path)
 
   def saved_model_inference(self, image_path_pattern, output_dir, **kwargs):
     """Perform inference for the given saved model."""
@@ -459,6 +462,7 @@ def main(argv):
       ckpt_path=FLAGS.ckpt_path,
       export_ckpt=FLAGS.export_ckpt,
       saved_model_dir=FLAGS.saved_model_dir,
+      tflite_path=FLAGS.tflite_path,
       batch_size=FLAGS.batch_size,
       hparams=FLAGS.hparams)
   inspector.run_model(