Browse Source

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
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
b4ac3df6ff
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 3 deletions
  1. +8
    -3
      export.py

+ 8
- 3
export.py View File

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
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}"
# 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,

Loading…
Cancel
Save