Unverified 提交 b4ac3df6 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

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
上级 00d7b978
...@@ -41,6 +41,7 @@ TensorFlow.js: ...@@ -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 ...@@ -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' ...@@ -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,
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论