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 <david.strahm@lambda-it.ch> * Update Dockerfile Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> --------- Signed-off-by: David Strahm <david.strahm@lambda-it.ch> Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>pull/10951/head
parent
80e5447390
commit
d389840f66
|
@ -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)
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue