Skip to content

Quantization-aware training using Tensorflow Object Detection API and compile Edge TPU model.

Notifications You must be signed in to change notification settings

habakan/object-detection-api-for-edgetpu

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 

Repository files navigation

Tensorflow Object Detection API for Edge TPU

Directory Structure

├── Dockerfile
├── README.md
├── config
│   ├── label_map.pbtxt
│   └── ssd_mobilenet_v1.config
├── out # output model file
└── utils
    ├── create_tf_record.py
    ├── create_tf_record_trainandval.py
    └── shell_script
        ├── convert_tfrecord.sh
        └── train_and_convert.sh

Requirements

Environment

  • Docker
  • nvidia-docker

File

  • Training dataset (Ex: PascalVOC dataset)
  • Model config (Ex: config/ssd_mobilenet_v1.config)
  • label_map.pbtxt (Ex: config/label_map.pbtxt)

Build Docker image

$ sudo dokcer build ./ -t tfodapi

Training pacalVOC dataset

Download dataset

$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
$ tar -xvzf VOCtrainval_11-May-2012.tar

Run Docker Conainer

$ sudo nvidia-docker run -it -v /path-to-dir/pascalVOC/train/VOCdevkit/VOC2012:/data \
		-v /path-to-repository/config:/config \
		-v /path-to-repository/out:/output \
		--name tfodapi tfodapi /bin/bash

Create tf-record in container

$ python object_detection/dataset_utils/create_tf_record.py \
		--annotations_dir=/data/Annotations --images_dir=/data/JPEGImages \
		--output_dir=/data/tfrecord --label_map_path=/config/label_map.pb

Training in container

$ cd /models/research
$ mkdir -p /output/train_model /output/tflite_model /output/inference_model
$ python object_detection/legacy/train.py --logtostderr --pipeline_config_path \
		/config/ssd_mobilenet_v1.config  --train_dir /output/train_model

[Option] Export Inference Graph

$ python object_detection/export_inference_graph.py --input_type image_tensor \
        --pipeline_config /output/train_model/pipeline.config_path \
        --trained_checkpoint_prefix /output/train_model/model.ckpt-200000 \
		--output_directory /output/inference_model

Build Edge TPU model in container

Convert tflite file

$ python object_detection/export_tflite_ssd_graph.py \
		--pipeline_config_path /output/train_model/pipeline.config \
		--trained_checkpoint_prefix /output/train_model/model.ckpt-200000 \
		--output_directory /output/tflite_model --add_postprocessing_op=true
$ tflite_convert --output_file /output/tflite_model/output_tflite_graph.tflite \
		--graph_def_file /output/tflite_model/tflite_graph.pb --inference_type=QUANTIZED_UINT8 \
    --input_arrays normalized_input_image_tensor \
    --output_arrays "TFLite_Detection_PostProcess,TFLite_Detection_PostProcess:1,TFLite_Detection_PostProcess:2,TFLite_Detection_PostProcess:3" \
    --mean_values 128 --std_dev_values 128 --input_shapes 1,300,300,3 \
    --change_concat_input_ranges false \
    --allow_nudging_weights_to_use_fast_gemm_kernel true --allow_custom_ops

Compile Edge TPU model

$ edgetpu_compiler /output/tflite_model/output_tflite_graph.tflite -o /output/tflite_model

Acknowledgement

About

Quantization-aware training using Tensorflow Object Detection API and compile Edge TPU model.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published