Skip to content

Commit

Permalink
Attempt edgetpu-compiler autoinstall (#6223)
Browse files Browse the repository at this point in the history
* Attempt `edgetpu-compiler` autoinstall

Attempt to install edgetpu-compiler dependency if missing on Linux.

* Update export.py

* Update export.py
  • Loading branch information
glenn-jocher committed Jan 6, 2022
1 parent b4ac3df commit f80c463
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -309,11 +309,16 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
# YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
try:
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)
ver = out.stdout.decode().split()[-1]
help_url = 'https://coral.ai/docs/edgetpu/compiler/'
assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'
if subprocess.run(cmd, shell=True).returncode != 0:
LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
'sudo apt-get update',
'sudo apt-get install edgetpu-compiler']:
subprocess.run(c, shell=True, check=True)
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]

LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
Expand Down

2 comments on commit f80c463

@positive666
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@glenn-jocher ,Hello! when judging the tensorrt version, there was a small error on line 177 that was a string type of ‘7’

@glenn-jocher
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@positive666 thanks for letting me know! I've fixed this in #6235 now :)

Please sign in to comment.