zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Add `edgetpu_compiler` checks (#6218) 

* Add `edgetpu_compiler` checks

* Update export.py

* Update export.py

* Update export.py

* Update export.py

* Update export.py

* Update export.py
This commit is contained in:
Glenn Jocher 2022-01-05 14:55:04 -08:00 committed by GitHub
parent 00d7b97869
commit b4ac3df6ff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
export.py
View File

@ -41,6 +41,7 @@ TensorFlow.js:
import argparse import argparse
import json import json
import os import os
import platform
import subprocess import subprocess
import sys import sys
import time import time
@ -307,11 +308,15 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
# YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
try: try:
cmd = 'edgetpu_compiler --version' # install https://coral.ai/docs/edgetpu/compiler/ cmd = 'edgetpu_compiler --version'
help = 'See https://coral.ai/docs/edgetpu/compiler/'
assert platform.system() == 'Linux', f'export only supported on Linux. {help}'
assert subprocess.run(cmd, shell=True).returncode == 0, f'export requires edgetpu-compiler. {help}'
out = subprocess.run(cmd, shell=True, capture_output=True, check=True) out = subprocess.run(cmd, shell=True, capture_output=True, check=True)
ver = out.stdout.decode().split()[-1] ver = out.stdout.decode().split()[-1]
LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...') LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
f = str(file).replace('.pt', '-int8_edgetpu.tflite') f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model
cmd = f"edgetpu_compiler -s {f_tfl}" cmd = f"edgetpu_compiler -s {f_tfl}"
@ -434,7 +439,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
# TensorFlow Exports # TensorFlow Exports
if any(tf_exports): if any(tf_exports):
pb, tflite, edgetpu, tfjs = tf_exports[1:] pb, tflite, edgetpu, tfjs = tf_exports[1:]
if (tflite or edgetpu) and int8: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707 if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow` check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs, model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,