From d389840f66bb95c150d8c0e4d97759b07d21e821 Mon Sep 17 00:00:00 2001 From: David Strahm Date: Fri, 10 Feb 2023 16:11:08 +0100 Subject: [PATCH] Allow int8 quantization for export_tfjs (#10948) * Allow int8 quantization for export_tfjs --int8 param currently has no effect on export_tfjs. With this change, ` python export.py --weights ../path/to/best.pt --include tfjs --int8` will add the --quantize_uint8 param to the tensorflowjs_converter script, greatly reducing model size for web usage. Signed-off-by: David Strahm * Update Dockerfile Signed-off-by: Glenn Jocher --------- Signed-off-by: David Strahm Signed-off-by: Glenn Jocher Co-authored-by: Glenn Jocher --- export.py | 8 +++++--- utils/docker/Dockerfile | 9 ++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/export.py b/export.py index 9ca3441..1bf0532 100644 --- a/export.py +++ b/export.py @@ -426,7 +426,7 @@ def export_edgetpu(file, prefix=colorstr('Edge TPU:')): @try_export -def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): +def export_tfjs(file, int8, prefix=colorstr('TensorFlow.js:')): # YOLOv5 TensorFlow.js export check_requirements('tensorflowjs') import tensorflowjs as tfjs @@ -436,7 +436,9 @@ def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): f_pb = file.with_suffix('.pb') # *.pb path f_json = f'{f}/model.json' # *.json path - cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \ + int8_export = ' --quantize_uint8 ' if int8 else '' + + cmd = f'tensorflowjs_converter --input_format=tf_frozen_model {int8_export}' \ f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}' subprocess.run(cmd.split()) @@ -588,7 +590,7 @@ def run( f[8], _ = export_edgetpu(file) add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs)) if tfjs: - f[9], _ = export_tfjs(file) + f[9], _ = export_tfjs(file, int8) if paddle: # PaddlePaddle f[10], _ = export_paddle(model, im, file, metadata) diff --git a/utils/docker/Dockerfile b/utils/docker/Dockerfile index 54927c0..cfe7b06 100644 --- a/utils/docker/Dockerfile +++ b/utils/docker/Dockerfile @@ -13,13 +13,12 @@ ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Aria ENV DEBIAN_FRONTEND noninteractive RUN apt update RUN TZ=Etc/UTC apt install -y tzdata -RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg +RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg \ + # Security updates + # https://security.snyk.io/vuln/SNYK-UBUNTU1804-OPENSSL-3314796 + openssl # RUN alias python=python3 -# Security updates -# https://security.snyk.io/vuln/SNYK-UBUNTU1804-OPENSSL-3314796 -RUN apt install --no-install-recommends -y openssl - # Create working directory RUN rm -rf /usr/src/app && mkdir -p /usr/src/app WORKDIR /usr/src/app