From c3a93d783d1a1e920d346f62b5de9f500e4540e4 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 12 Sep 2021 15:52:24 +0200 Subject: [PATCH] Add TensorFlow formats to `export.py` (#4479) * Initial commit * Remove unused export_torchscript return * ROOT variable * Add prefix to fcn arg * fix ROOT * check_yaml into run() * interim fixes * imgsz=(320, 320) * Hardcode tf_raw_resize False * Finish opt elimination * Update representative_dataset_gen() * Update export.py with TF methods * SiLU and GraphDef fixes * file_size() directory handling feature * export fixes * add lambda: to representative_dataset * Detect training False default * Fuse false for TF models * Embed agnostic NMS arguments * Remove lambda * TensorFlow.js export success * Add pb to Usage * Add *_tfjs_model/ to ignore files * prepend YOLOv5 to function headers * Remove end --- comments * parameterize tfjs export pb file * update run() data default /ROOT * update --include help * update imports * return ct_model * Consolidate TFLite export * pb prerequisite to tfjs * TF modules CamelCase * Remove exports from tf.py and cleanup * pass agnostic NMS arguments * CI * CI * ignore *_web_model/ * Add tensorflow to CI dependencies * CI tensorflow-cpu * Update requirements.txt * Remove tensorflow check_requirement * CI coreml tfjs * export only onnx torchscript * reorder exports torchscript first --- .dockerignore | 1 + .github/workflows/ci-testing.yml | 7 +- .gitignore | 1 + detect.py | 2 +- export.py | 219 +++++++++++++--- models/tf.py | 433 ++++++++++++------------------- requirements.txt | 20 +- utils/general.py | 12 +- 8 files changed, 366 insertions(+), 329 deletions(-) diff --git a/.dockerignore b/.dockerignore index 8d60b462e7d1..6c2f2b9b7725 100644 --- a/.dockerignore +++ b/.dockerignore @@ -22,6 +22,7 @@ data/samples/* **/*.h5 **/*.pb *_saved_model/ +*_web_model/ # Below Copied From .gitignore ----------------------------------------------------------------------------------------- # Below Copied From .gitignore ----------------------------------------------------------------------------------------- diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml index ecd6f9bbd625..54b230a13e6b 100644 --- a/.github/workflows/ci-testing.yml +++ b/.github/workflows/ci-testing.yml @@ -48,7 +48,7 @@ jobs: run: | python -m pip install --upgrade pip pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html - pip install -q onnx onnx-simplifier coremltools # for export + pip install -q onnx tensorflow-cpu # for export python --version pip --version pip list @@ -75,6 +75,7 @@ jobs: python val.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di python hubconf.py # hub - python models/yolo.py --cfg ${{ matrix.model }}.yaml # inspect - python export.py --img 128 --batch 1 --weights ${{ matrix.model }}.pt --include onnx torchscript # export + python models/yolo.py --cfg ${{ matrix.model }}.yaml # build PyTorch model + python models/tf.py --weights ${{ matrix.model }}.pt # build TensorFlow model + python export.py --img 128 --batch 1 --weights ${{ matrix.model }}.pt --include torchscript onnx # export shell: bash diff --git a/.gitignore b/.gitignore index f8a2437973f0..375b71807588 100755 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,7 @@ VOC/ *.tflite *.h5 *_saved_model/ +*_web_model/ darknet53.conv.74 yolov3-tiny.conv.15 diff --git a/detect.py b/detect.py index b6597c1662f9..ef7458d52db3 100644 --- a/detect.py +++ b/detect.py @@ -253,7 +253,7 @@ def wrap_frozen_graph(gd, inputs, outputs): def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') + parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model path(s)') parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam') parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold') diff --git a/export.py b/export.py index 935bdb40bc9b..8d6805893d1e 100644 --- a/export.py +++ b/export.py @@ -1,12 +1,28 @@ # YOLOv5 🚀 by Ultralytics, GPL-3.0 license """ -Export a PyTorch model to TorchScript, ONNX, CoreML formats +Export a YOLOv5 PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats +TensorFlow exports authored by https://github.com/zldrobit Usage: - $ python path/to/export.py --weights yolov5s.pt --img 640 --batch 1 + $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs + +Inference: + $ python path/to/detect.py --weights yolov5s.pt + yolov5s.onnx (must export with --dynamic) + yolov5s_saved_model + yolov5s.pb + yolov5s.tflite + +TensorFlow.js: + $ # Edit yolov5s_web_model/model.json to sort Identity* in ascending order + $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example + $ npm install + $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model + $ npm start """ import argparse +import subprocess import sys import time from pathlib import Path @@ -16,40 +32,42 @@ from torch.utils.mobile_optimizer import optimize_for_mobile FILE = Path(__file__).resolve() -sys.path.append(FILE.parents[0].as_posix()) # add yolov5/ to path +ROOT = FILE.parents[0] # yolov5/ dir +sys.path.append(ROOT.as_posix()) # add yolov5/ to path from models.common import Conv -from models.yolo import Detect from models.experimental import attempt_load -from utils.activations import Hardswish, SiLU -from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging +from models.yolo import Detect +from utils.activations import SiLU +from utils.datasets import LoadImages +from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, set_logging from utils.torch_utils import select_device -def export_torchscript(model, img, file, optimize): - # TorchScript model export - prefix = colorstr('TorchScript:') +def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')): + # YOLOv5 TorchScript model export try: print(f'\n{prefix} starting export with torch {torch.__version__}...') f = file.with_suffix('.torchscript.pt') - ts = torch.jit.trace(model, img, strict=False) + + ts = torch.jit.trace(model, im, strict=False) (optimize_for_mobile(ts) if optimize else ts).save(f) + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') - return ts except Exception as e: print(f'{prefix} export failure: {e}') -def export_onnx(model, img, file, opset, train, dynamic, simplify): - # ONNX model export - prefix = colorstr('ONNX:') +def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')): + # YOLOv5 ONNX export try: check_requirements(('onnx',)) import onnx print(f'\n{prefix} starting export with onnx {onnx.__version__}...') f = file.with_suffix('.onnx') - torch.onnx.export(model, img, f, verbose=False, opset_version=opset, + + torch.onnx.export(model, im, f, verbose=False, opset_version=opset, training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, do_constant_folding=not train, input_names=['images'], @@ -73,7 +91,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify): model_onnx, check = onnxsim.simplify( model_onnx, dynamic_input_shape=dynamic, - input_shapes={'images': list(img.shape)} if dynamic else None) + input_shapes={'images': list(im.shape)} if dynamic else None) assert check, 'assert check failed' onnx.save(model_onnx, f) except Exception as e: @@ -84,26 +102,131 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify): print(f'{prefix} export failure: {e}') -def export_coreml(model, img, file): - # CoreML model export - prefix = colorstr('CoreML:') +def export_coreml(model, im, file, prefix=colorstr('CoreML:')): + # YOLOv5 CoreML export + ct_model = None try: check_requirements(('coremltools',)) import coremltools as ct print(f'\n{prefix} starting export with coremltools {ct.__version__}...') f = file.with_suffix('.mlmodel') + model.train() # CoreML exports should be placed in model.train() mode - ts = torch.jit.trace(model, img, strict=False) # TorchScript model - model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) - model.save(f) + ts = torch.jit.trace(model, im, strict=False) # TorchScript model + ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255.0, bias=[0, 0, 0])]) + ct_model.save(f) + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') except Exception as e: print(f'\n{prefix} export failure: {e}') + return ct_model + + +def export_saved_model(model, im, file, dynamic, + tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, + conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')): + # YOLOv5 TensorFlow saved_model export + keras_model = None + try: + import tensorflow as tf + from tensorflow import keras + from models.tf import TFModel, TFDetect + + print(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + f = str(file).replace('.pt', '_saved_model') + batch_size, ch, *imgsz = list(im.shape) # BCHW -def run(weights='./yolov5s.pt', # weights path - img_size=(640, 640), # image (height, width) + tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) + im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow + y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size) + outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + keras_model = keras.Model(inputs=inputs, outputs=outputs) + keras_model.summary() + keras_model.save(f, save_format='tf') + + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + except Exception as e: + print(f'\n{prefix} export failure: {e}') + + return keras_model + + +def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): + # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow + try: + import tensorflow as tf + from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 + + print(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + f = file.with_suffix('.pb') + + m = tf.function(lambda x: keras_model(x)) # full model + m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) + frozen_func = convert_variables_to_constants_v2(m) + frozen_func.graph.as_graph_def() + tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) + + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + except Exception as e: + print(f'\n{prefix} export failure: {e}') + + +def export_tflite(keras_model, im, file, tfl_int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')): + # YOLOv5 TensorFlow Lite export + try: + import tensorflow as tf + from models.tf import representative_dataset_gen + + print(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + batch_size, ch, *imgsz = list(im.shape) # BCHW + f = file.with_suffix('.tflite') + + converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) + converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] + converter.optimizations = [tf.lite.Optimize.DEFAULT] + if tfl_int8: + dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data + converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib) + converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] + converter.inference_input_type = tf.uint8 # or tf.int8 + converter.inference_output_type = tf.uint8 # or tf.int8 + converter.experimental_new_quantizer = False + f = str(file).replace('.pt', '-int8.tflite') + + tflite_model = converter.convert() + open(f, "wb").write(tflite_model) + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + + except Exception as e: + print(f'\n{prefix} export failure: {e}') + + +def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): + # YOLOv5 TensorFlow.js export + try: + check_requirements(('tensorflowjs',)) + import tensorflowjs as tfjs + + print(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') + f = str(file).replace('.pt', '_web_model') # js dir + f_pb = file.with_suffix('.pb') # *.pb path + + cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \ + f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}" + subprocess.run(cmd, shell=True) + + print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + except Exception as e: + print(f'\n{prefix} export failure: {e}') + + +@torch.no_grad() +def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' + weights=ROOT / 'yolov5s.pt', # weights path + imgsz=(640, 640), # image (height, width) batch_size=1, # batch size device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu include=('torchscript', 'onnx', 'coreml'), # include formats @@ -117,29 +240,28 @@ def run(weights='./yolov5s.pt', # weights path ): t = time.time() include = [x.lower() for x in include] - img_size *= 2 if len(img_size) == 1 else 1 # expand + tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports + imgsz *= 2 if len(imgsz) == 1 else 1 # expand file = Path(weights) # Load PyTorch model device = select_device(device) assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0' - model = attempt_load(weights, map_location=device) # load FP32 model - names = model.names + model = attempt_load(weights, map_location=device, inplace=True, fuse=not any(tf_exports)) # load FP32 model + nc, names = model.nc, model.names # number of classes, class names # Input gs = int(max(model.stride)) # grid size (max stride) - img_size = [check_img_size(x, gs) for x in img_size] # verify img_size are gs-multiples - img = torch.zeros(batch_size, 3, *img_size).to(device) # image size(1,3,320,192) iDetection + imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples + im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection # Update model if half: - img, model = img.half(), model.half() # to FP16 + im, model = im.half(), model.half() # to FP16 model.train() if train else model.eval() # training mode = no Detect() layer grid construction for k, m in model.named_modules(): if isinstance(m, Conv): # assign export-friendly activations - if isinstance(m.act, nn.Hardswish): - m.act = Hardswish() - elif isinstance(m.act, nn.SiLU): + if isinstance(m.act, nn.SiLU): m.act = SiLU() elif isinstance(m, Detect): m.inplace = inplace @@ -147,16 +269,28 @@ def run(weights='./yolov5s.pt', # weights path # m.forward = m.forward_export # assign forward (optional) for _ in range(2): - y = model(img) # dry runs + y = model(im) # dry runs print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)") # Exports if 'torchscript' in include: - export_torchscript(model, img, file, optimize) + export_torchscript(model, im, file, optimize) if 'onnx' in include: - export_onnx(model, img, file, opset, train, dynamic, simplify) + export_onnx(model, im, file, opset, train, dynamic, simplify) if 'coreml' in include: - export_coreml(model, img, file) + export_coreml(model, im, file) + + # TensorFlow Exports + if any(tf_exports): + pb, tflite, tfjs = tf_exports[1:] + 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=tfjs, agnostic_nms=tfjs) # keras model + if pb or tfjs: # pb prerequisite to tfjs + export_pb(model, im, file) + if tflite: + export_tflite(model, im, file, tfl_int8=False, data=data, ncalib=100) + if tfjs: + export_tfjs(model, im, file) # Finish print(f'\nExport complete ({time.time() - t:.2f}s)' @@ -166,18 +300,21 @@ def run(weights='./yolov5s.pt', # weights path def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') - parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (height, width)') + parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') + parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') + parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)') parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats') parser.add_argument('--half', action='store_true', help='FP16 half-precision export') parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True') parser.add_argument('--train', action='store_true', help='model.train() mode') parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') - parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes') + parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version') + parser.add_argument('--include', nargs='+', + default=['torchscript', 'onnx'], + help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)') opt = parser.parse_args() return opt diff --git a/models/tf.py b/models/tf.py index d6d0f26210b2..621236240f10 100644 --- a/models/tf.py +++ b/models/tf.py @@ -1,67 +1,44 @@ # YOLOv5 🚀 by Ultralytics, GPL-3.0 license """ -TensorFlow/Keras and TFLite versions of YOLOv5 +TensorFlow, Keras and TFLite versions of YOLOv5 Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127 Usage: - $ python models/tf.py --weights yolov5s.pt --cfg yolov5s.yaml - -Export int8 TFLite models: - $ python models/tf.py --weights yolov5s.pt --cfg models/yolov5s.yaml --tfl-int8 \ - --source path/to/images/ --ncalib 100 - -Detection: - $ python detect.py --weights yolov5s.pb --img 320 - $ python detect.py --weights yolov5s_saved_model --img 320 - $ python detect.py --weights yolov5s-fp16.tflite --img 320 - $ python detect.py --weights yolov5s-int8.tflite --img 320 --tfl-int8 - -For TensorFlow.js: - $ python models/tf.py --weights yolov5s.pt --cfg models/yolov5s.yaml --img 320 --tf-nms --agnostic-nms - $ pip install tensorflowjs - $ tensorflowjs_converter \ - --input_format=tf_frozen_model \ - --output_node_names='Identity,Identity_1,Identity_2,Identity_3' \ - yolov5s.pb \ - web_model - $ # Edit web_model/model.json to sort Identity* in ascending order - $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example - $ npm install - $ ln -s ../../yolov5/web_model public/web_model - $ npm start + $ python models/tf.py --weights yolov5s.pt + +Export: + $ python path/to/export.py --weights yolov5s.pt --include saved_model pb tflite tfjs """ import argparse import logging -import os import sys -import traceback from copy import deepcopy from pathlib import Path -sys.path.append('./') # to run '$ python *.py' files in subdirectories +FILE = Path(__file__).resolve() +ROOT = FILE.parents[1] # yolov5/ dir +sys.path.append(ROOT.as_posix()) # add yolov5/ to path import numpy as np import tensorflow as tf import torch import torch.nn as nn -import yaml from tensorflow import keras -from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, autopad, C3 from models.experimental import MixConv2d, CrossConv, attempt_load from models.yolo import Detect -from utils.datasets import LoadImages -from utils.general import check_dataset, check_yaml, make_divisible +from utils.general import colorstr, make_divisible, set_logging +from utils.activations import SiLU -logger = logging.getLogger(__name__) +LOGGER = logging.getLogger(__name__) -class tf_BN(keras.layers.Layer): +class TFBN(keras.layers.Layer): # TensorFlow BatchNormalization wrapper def __init__(self, w=None): - super(tf_BN, self).__init__() + super(TFBN, self).__init__() self.bn = keras.layers.BatchNormalization( beta_initializer=keras.initializers.Constant(w.bias.numpy()), gamma_initializer=keras.initializers.Constant(w.weight.numpy()), @@ -73,20 +50,20 @@ def call(self, inputs): return self.bn(inputs) -class tf_Pad(keras.layers.Layer): +class TFPad(keras.layers.Layer): def __init__(self, pad): - super(tf_Pad, self).__init__() + super(TFPad, self).__init__() self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]]) def call(self, inputs): return tf.pad(inputs, self.pad, mode='constant', constant_values=0) -class tf_Conv(keras.layers.Layer): +class TFConv(keras.layers.Layer): # Standard convolution def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): # ch_in, ch_out, weights, kernel, stride, padding, groups - super(tf_Conv, self).__init__() + super(TFConv, self).__init__() assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" assert isinstance(k, int), "Convolution with multiple kernels are not allowed." # TensorFlow convolution padding is inconsistent with PyTorch (e.g. k=3 s=2 'SAME' padding) @@ -95,27 +72,29 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): conv = keras.layers.Conv2D( c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False, kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy())) - self.conv = conv if s == 1 else keras.Sequential([tf_Pad(autopad(k, p)), conv]) - self.bn = tf_BN(w.bn) if hasattr(w, 'bn') else tf.identity + self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv]) + self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity # YOLOv5 activations if isinstance(w.act, nn.LeakyReLU): self.act = (lambda x: keras.activations.relu(x, alpha=0.1)) if act else tf.identity elif isinstance(w.act, nn.Hardswish): self.act = (lambda x: x * tf.nn.relu6(x + 3) * 0.166666667) if act else tf.identity - elif isinstance(w.act, nn.SiLU): + elif isinstance(w.act, (nn.SiLU, SiLU)): self.act = (lambda x: keras.activations.swish(x)) if act else tf.identity + else: + raise Exception(f'no matching TensorFlow activation found for {w.act}') def call(self, inputs): return self.act(self.bn(self.conv(inputs))) -class tf_Focus(keras.layers.Layer): +class TFFocus(keras.layers.Layer): # Focus wh information into c-space def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): # ch_in, ch_out, kernel, stride, padding, groups - super(tf_Focus, self).__init__() - self.conv = tf_Conv(c1 * 4, c2, k, s, p, g, act, w.conv) + super(TFFocus, self).__init__() + self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv) def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c) # inputs = inputs / 255. # normalize 0-255 to 0-1 @@ -125,23 +104,23 @@ def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c) inputs[:, 1::2, 1::2, :]], 3)) -class tf_Bottleneck(keras.layers.Layer): +class TFBottleneck(keras.layers.Layer): # Standard bottleneck def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None): # ch_in, ch_out, shortcut, groups, expansion - super(tf_Bottleneck, self).__init__() + super(TFBottleneck, self).__init__() c_ = int(c2 * e) # hidden channels - self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1) - self.cv2 = tf_Conv(c_, c2, 3, 1, g=g, w=w.cv2) + self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) + self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2) self.add = shortcut and c1 == c2 def call(self, inputs): return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs)) -class tf_Conv2d(keras.layers.Layer): +class TFConv2d(keras.layers.Layer): # Substitution for PyTorch nn.Conv2D def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None): - super(tf_Conv2d, self).__init__() + super(TFConv2d, self).__init__() assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" self.conv = keras.layers.Conv2D( c2, k, s, 'VALID', use_bias=bias, @@ -152,19 +131,19 @@ def call(self, inputs): return self.conv(inputs) -class tf_BottleneckCSP(keras.layers.Layer): +class TFBottleneckCSP(keras.layers.Layer): # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): # ch_in, ch_out, number, shortcut, groups, expansion - super(tf_BottleneckCSP, self).__init__() + super(TFBottleneckCSP, self).__init__() c_ = int(c2 * e) # hidden channels - self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1) - self.cv2 = tf_Conv2d(c1, c_, 1, 1, bias=False, w=w.cv2) - self.cv3 = tf_Conv2d(c_, c_, 1, 1, bias=False, w=w.cv3) - self.cv4 = tf_Conv(2 * c_, c2, 1, 1, w=w.cv4) - self.bn = tf_BN(w.bn) + self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) + self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2) + self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3) + self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4) + self.bn = TFBN(w.bn) self.act = lambda x: keras.activations.relu(x, alpha=0.1) - self.m = keras.Sequential([tf_Bottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) + self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) def call(self, inputs): y1 = self.cv3(self.m(self.cv1(inputs))) @@ -172,28 +151,28 @@ def call(self, inputs): return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3)))) -class tf_C3(keras.layers.Layer): +class TFC3(keras.layers.Layer): # CSP Bottleneck with 3 convolutions def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): # ch_in, ch_out, number, shortcut, groups, expansion - super(tf_C3, self).__init__() + super(TFC3, self).__init__() c_ = int(c2 * e) # hidden channels - self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1) - self.cv2 = tf_Conv(c1, c_, 1, 1, w=w.cv2) - self.cv3 = tf_Conv(2 * c_, c2, 1, 1, w=w.cv3) - self.m = keras.Sequential([tf_Bottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) + self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) + self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2) + self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3) + self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)]) def call(self, inputs): return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3)) -class tf_SPP(keras.layers.Layer): +class TFSPP(keras.layers.Layer): # Spatial pyramid pooling layer used in YOLOv3-SPP def __init__(self, c1, c2, k=(5, 9, 13), w=None): - super(tf_SPP, self).__init__() + super(TFSPP, self).__init__() c_ = c1 // 2 # hidden channels - self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1) - self.cv2 = tf_Conv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2) + self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) + self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2) self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k] def call(self, inputs): @@ -201,9 +180,9 @@ def call(self, inputs): return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3)) -class tf_Detect(keras.layers.Layer): - def __init__(self, nc=80, anchors=(), ch=(), w=None): # detection layer - super(tf_Detect, self).__init__() +class TFDetect(keras.layers.Layer): + def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None): # detection layer + super(TFDetect, self).__init__() self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32) self.nc = nc # number of classes self.no = nc + 5 # number of outputs per anchor @@ -213,22 +192,20 @@ def __init__(self, nc=80, anchors=(), ch=(), w=None): # detection layer self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32) self.anchor_grid = tf.reshape(tf.convert_to_tensor(w.anchor_grid.numpy(), dtype=tf.float32), [self.nl, 1, -1, 1, 2]) - self.m = [tf_Conv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)] - self.export = False # onnx export - self.training = True # set to False after building model + self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)] + self.training = False # set to False after building model + self.imgsz = imgsz for i in range(self.nl): - ny, nx = opt.img_size[0] // self.stride[i], opt.img_size[1] // self.stride[i] + ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i] self.grid[i] = self._make_grid(nx, ny) def call(self, inputs): - # x = x.copy() # for profiling z = [] # inference output - self.training |= self.export x = [] for i in range(self.nl): x.append(self.m[i](inputs[i])) # x(bs,20,20,255) to x(bs,3,20,20,85) - ny, nx = opt.img_size[0] // self.stride[i], opt.img_size[1] // self.stride[i] + ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i] x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na, self.no]), [0, 2, 1, 3]) if not self.training: # inference @@ -236,8 +213,8 @@ def call(self, inputs): xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # Normalize xywh to 0-1 to reduce calibration error - xy /= tf.constant([[opt.img_size[1], opt.img_size[0]]], dtype=tf.float32) - wh /= tf.constant([[opt.img_size[1], opt.img_size[0]]], dtype=tf.float32) + xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) + wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32) y = tf.concat([xy, wh, y[..., 4:]], -1) z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no])) @@ -251,25 +228,23 @@ def _make_grid(nx=20, ny=20): return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32) -class tf_Upsample(keras.layers.Layer): - def __init__(self, size, scale_factor, mode, w=None): - super(tf_Upsample, self).__init__() +class TFUpsample(keras.layers.Layer): + def __init__(self, size, scale_factor, mode, w=None): # warning: all arguments needed including 'w' + super(TFUpsample, self).__init__() assert scale_factor == 2, "scale_factor must be 2" + self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode) # self.upsample = keras.layers.UpSampling2D(size=scale_factor, interpolation=mode) - if opt.tf_raw_resize: - # with default arguments: align_corners=False, half_pixel_centers=False - self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x, - size=(x.shape[1] * 2, x.shape[2] * 2)) - else: - self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode) + # with default arguments: align_corners=False, half_pixel_centers=False + # self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x, + # size=(x.shape[1] * 2, x.shape[2] * 2)) def call(self, inputs): return self.upsample(inputs) -class tf_Concat(keras.layers.Layer): +class TFConcat(keras.layers.Layer): def __init__(self, dimension=1, w=None): - super(tf_Concat, self).__init__() + super(TFConcat, self).__init__() assert dimension == 1, "convert only NCHW to NHWC concat" self.d = 3 @@ -277,8 +252,8 @@ def call(self, inputs): return tf.concat(inputs, self.d) -def parse_model(d, ch, model): # model_dict, input_channels(3) - logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) +def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) + LOGGER.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors no = na * (nc + 5) # number of outputs = anchors * (classes + 5) @@ -310,10 +285,11 @@ def parse_model(d, ch, model): # model_dict, input_channels(3) args.append([ch[x + 1] for x in f]) if isinstance(args[1], int): # number of anchors args[1] = [list(range(args[1] * 2))] * len(f) + args.append(imgsz) else: c2 = ch[f] - tf_m = eval('tf_' + m_str.replace('nn.', '')) + tf_m = eval('TF' + m_str.replace('nn.', '')) m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \ else tf_m(*args, w=model.model[i]) # module @@ -321,16 +297,16 @@ def parse_model(d, ch, model): # model_dict, input_channels(3) t = str(m)[8:-2].replace('__main__.', '') # module type np = sum([x.numel() for x in torch_m_.parameters()]) # number params m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params - logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print + LOGGER.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist layers.append(m_) ch.append(c2) return keras.Sequential(layers), sorted(save) -class tf_Model(): - def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None): # model, input channels, number of classes - super(tf_Model, self).__init__() +class TFModel: + def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)): # model, channels, classes + super(TFModel, self).__init__() if isinstance(cfg, dict): self.yaml = cfg # model dict else: # is *.yaml @@ -343,9 +319,10 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None): # model, inp if nc and nc != self.yaml['nc']: print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) self.yaml['nc'] = nc # override yaml value - self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model) # model, savelist, ch_out + self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz) - def predict(self, inputs, profile=False): + def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, + conf_thres=0.25): y = [] # outputs x = inputs for i, m in enumerate(self.model.layers): @@ -356,18 +333,18 @@ def predict(self, inputs, profile=False): y.append(x if m.i in self.savelist else None) # save output # Add TensorFlow NMS - if opt.tf_nms: - boxes = xywh2xyxy(x[0][..., :4]) + if tf_nms: + boxes = self._xywh2xyxy(x[0][..., :4]) probs = x[0][:, :, 4:5] classes = x[0][:, :, 5:] scores = probs * classes - if opt.agnostic_nms: - nms = agnostic_nms_layer()((boxes, classes, scores)) + if agnostic_nms: + nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres) return nms, x[1] else: boxes = tf.expand_dims(boxes, 2) nms = tf.image.combined_non_max_suppression( - boxes, scores, opt.topk_per_class, opt.topk_all, opt.iou_thres, opt.score_thres, clip_boxes=False) + boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False) return nms, x[1] return x[0] # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...] @@ -377,182 +354,94 @@ def predict(self, inputs, profile=False): # cls = tf.reshape(tf.cast(tf.argmax(x[..., 5:], axis=1), tf.float32), (-1, 1)) # x(6300,1) classes # return tf.concat([conf, cls, xywh], 1) + @staticmethod + def _xywh2xyxy(xywh): + # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right + x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1) + return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1) + -class agnostic_nms_layer(keras.layers.Layer): - # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450 - def call(self, input): - return tf.map_fn(agnostic_nms, input, +class AgnosticNMS(keras.layers.Layer): + # TF Agnostic NMS + def call(self, input, topk_all, iou_thres, conf_thres): + # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450 + return tf.map_fn(self._nms, input, fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32), name='agnostic_nms') - -def agnostic_nms(x): - boxes, classes, scores = x - class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32) - scores_inp = tf.reduce_max(scores, -1) - selected_inds = tf.image.non_max_suppression( - boxes, scores_inp, max_output_size=opt.topk_all, iou_threshold=opt.iou_thres, score_threshold=opt.score_thres) - selected_boxes = tf.gather(boxes, selected_inds) - padded_boxes = tf.pad(selected_boxes, - paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]], [0, 0]], - mode="CONSTANT", constant_values=0.0) - selected_scores = tf.gather(scores_inp, selected_inds) - padded_scores = tf.pad(selected_scores, - paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]]], - mode="CONSTANT", constant_values=-1.0) - selected_classes = tf.gather(class_inds, selected_inds) - padded_classes = tf.pad(selected_classes, - paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]]], - mode="CONSTANT", constant_values=-1.0) - valid_detections = tf.shape(selected_inds)[0] - return padded_boxes, padded_scores, padded_classes, valid_detections - - -def xywh2xyxy(xywh): - # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right - x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1) - return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1) - - -def representative_dataset_gen(): - # Representative dataset for use with converter.representative_dataset - n = 0 - for path, img, im0s, vid_cap in dataset: - # Get sample input data as a numpy array in a method of your choosing. - n += 1 + @staticmethod + def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25): # agnostic NMS + boxes, classes, scores = x + class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32) + scores_inp = tf.reduce_max(scores, -1) + selected_inds = tf.image.non_max_suppression( + boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres) + selected_boxes = tf.gather(boxes, selected_inds) + padded_boxes = tf.pad(selected_boxes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], + mode="CONSTANT", constant_values=0.0) + selected_scores = tf.gather(scores_inp, selected_inds) + padded_scores = tf.pad(selected_scores, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", constant_values=-1.0) + selected_classes = tf.gather(class_inds, selected_inds) + padded_classes = tf.pad(selected_classes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", constant_values=-1.0) + valid_detections = tf.shape(selected_inds)[0] + return padded_boxes, padded_scores, padded_classes, valid_detections + + +def representative_dataset_gen(dataset, ncalib=100): + # Representative dataset generator for use with converter.representative_dataset, returns a generator of np arrays + for n, (path, img, im0s, vid_cap) in enumerate(dataset): input = np.transpose(img, [1, 2, 0]) input = np.expand_dims(input, axis=0).astype(np.float32) input /= 255.0 yield [input] - if n >= opt.ncalib: + if n >= ncalib: break -if __name__ == "__main__": +def run(weights=ROOT / 'yolov5s.pt', # weights path + imgsz=(640, 640), # inference size h,w + batch_size=1, # batch size + dynamic=False, # dynamic batch size + ): + # PyTorch model + im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image + model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False) + y = model(im) # inference + model.info() + + # TensorFlow model + im = tf.zeros((batch_size, *imgsz, 3)) # BHWC image + tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) + y = tf_model.predict(im) # inference + + # Keras model + im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size) + keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im)) + keras_model.summary() + + +def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='cfg path') - parser.add_argument('--weights', type=str, default='yolov5s.pt', help='weights path') - parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size') # height, width + parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') + parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--dynamic-batch-size', action='store_true', help='dynamic batch size') - parser.add_argument('--source', type=str, default='../data/coco128.yaml', help='dir of images or data.yaml file') - parser.add_argument('--ncalib', type=int, default=100, help='number of calibration images') - parser.add_argument('--tfl-int8', action='store_true', dest='tfl_int8', help='export TFLite int8 model') - parser.add_argument('--tf-nms', action='store_true', dest='tf_nms', help='TF NMS (without TFLite export)') - parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') - parser.add_argument('--tf-raw-resize', action='store_true', dest='tf_raw_resize', - help='use tf.raw_ops.ResizeNearestNeighbor for resize') - parser.add_argument('--topk-per-class', type=int, default=100, help='topk per class to keep in NMS') - parser.add_argument('--topk-all', type=int, default=100, help='topk for all classes to keep in NMS') - parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') - parser.add_argument('--score-thres', type=float, default=0.4, help='score threshold for NMS') + parser.add_argument('--dynamic', action='store_true', help='dynamic batch size') opt = parser.parse_args() - opt.cfg = check_yaml(opt.cfg) # check YAML - opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand - print(opt) - - # Input - img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection - - # Load PyTorch model - model = attempt_load(opt.weights, map_location=torch.device('cpu'), inplace=True, fuse=False) - model.model[-1].export = False # set Detect() layer export=True - y = model(img) # dry run - nc = y[0].shape[-1] - 5 - - # TensorFlow saved_model export - try: - print('\nStarting TensorFlow saved_model export with TensorFlow %s...' % tf.__version__) - tf_model = tf_Model(opt.cfg, model=model, nc=nc) - img = tf.zeros((opt.batch_size, *opt.img_size, 3)) # NHWC Input for TensorFlow - - m = tf_model.model.layers[-1] - assert isinstance(m, tf_Detect), "the last layer must be Detect" - m.training = False - y = tf_model.predict(img) - - inputs = keras.Input(shape=(*opt.img_size, 3), batch_size=None if opt.dynamic_batch_size else opt.batch_size) - keras_model = keras.Model(inputs=inputs, outputs=tf_model.predict(inputs)) - keras_model.summary() - path = opt.weights.replace('.pt', '_saved_model') # filename - keras_model.save(path, save_format='tf') - print('TensorFlow saved_model export success, saved as %s' % path) - except Exception as e: - print('TensorFlow saved_model export failure: %s' % e) - traceback.print_exc(file=sys.stdout) - - # TensorFlow GraphDef export - try: - print('\nStarting TensorFlow GraphDef export with TensorFlow %s...' % tf.__version__) - - # https://github.com/leimao/Frozen_Graph_TensorFlow - full_model = tf.function(lambda x: keras_model(x)) - full_model = full_model.get_concrete_function( - tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) - - frozen_func = convert_variables_to_constants_v2(full_model) - frozen_func.graph.as_graph_def() - f = opt.weights.replace('.pt', '.pb') # filename - tf.io.write_graph(graph_or_graph_def=frozen_func.graph, - logdir=os.path.dirname(f), - name=os.path.basename(f), - as_text=False) - - print('TensorFlow GraphDef export success, saved as %s' % f) - except Exception as e: - print('TensorFlow GraphDef export failure: %s' % e) - traceback.print_exc(file=sys.stdout) - - # TFLite model export - if not opt.tf_nms: - try: - print('\nStarting TFLite export with TensorFlow %s...' % tf.__version__) - - # fp32 TFLite model export --------------------------------------------------------------------------------- - # converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) - # converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] - # converter.allow_custom_ops = False - # converter.experimental_new_converter = True - # tflite_model = converter.convert() - # f = opt.weights.replace('.pt', '.tflite') # filename - # open(f, "wb").write(tflite_model) - - # fp16 TFLite model export --------------------------------------------------------------------------------- - converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) - converter.optimizations = [tf.lite.Optimize.DEFAULT] - # converter.representative_dataset = representative_dataset_gen - # converter.target_spec.supported_types = [tf.float16] - converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] - converter.allow_custom_ops = False - converter.experimental_new_converter = True - tflite_model = converter.convert() - f = opt.weights.replace('.pt', '-fp16.tflite') # filename - open(f, "wb").write(tflite_model) - print('\nTFLite export success, saved as %s' % f) - - # int8 TFLite model export --------------------------------------------------------------------------------- - if opt.tfl_int8: - # Representative Dataset - if opt.source.endswith('.yaml'): - with open(check_yaml(opt.source)) as f: - data = yaml.load(f, Loader=yaml.FullLoader) # data dict - check_dataset(data) # check - opt.source = data['train'] - dataset = LoadImages(opt.source, img_size=opt.img_size, auto=False) - converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) - converter.optimizations = [tf.lite.Optimize.DEFAULT] - converter.representative_dataset = representative_dataset_gen - converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] - converter.inference_input_type = tf.uint8 # or tf.int8 - converter.inference_output_type = tf.uint8 # or tf.int8 - converter.allow_custom_ops = False - converter.experimental_new_converter = True - converter.experimental_new_quantizer = False - tflite_model = converter.convert() - f = opt.weights.replace('.pt', '-int8.tflite') # filename - open(f, "wb").write(tflite_model) - print('\nTFLite (int8) export success, saved as %s' % f) - - except Exception as e: - print('\nTFLite export failure: %s' % e) - traceback.print_exc(file=sys.stdout) + opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand + return opt + + +def main(opt): + set_logging() + print(colorstr('tf.py: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items())) + run(**vars(opt)) + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/requirements.txt b/requirements.txt index 2ad65ba53e29..b84b353f75f3 100755 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ # pip install -r requirements.txt -# base ---------------------------------------- +# Base ---------------------------------------- matplotlib>=3.2.2 numpy>=1.18.5 opencv-python>=4.1.2 @@ -11,21 +11,23 @@ torch>=1.7.0 torchvision>=0.8.1 tqdm>=4.41.0 -# logging ------------------------------------- +# Logging ------------------------------------- tensorboard>=2.4.1 # wandb -# plotting ------------------------------------ +# Plotting ------------------------------------ seaborn>=0.11.0 pandas -# export -------------------------------------- -# coremltools>=4.1 -# onnx>=1.9.0 -# scikit-learn==0.19.2 # for coreml quantization -# tensorflow==2.4.1 # for TFLite export +# Export -------------------------------------- +# coremltools>=4.1 # CoreML export +# onnx>=1.9.0 # ONNX export +# onnx-simplifier>=0.3.6 # ONNX simplifier +# scikit-learn==0.19.2 # CoreML quantization +# tensorflow>=2.4.1 # TFLite export +# tensorflowjs>=3.9.0 # TF.js export -# extras -------------------------------------- +# Extras -------------------------------------- # Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172 # pycocotools>=2.0 # COCO mAP # albumentations>=1.0.3 diff --git a/utils/general.py b/utils/general.py index 5c3d8d117dc3..7a80b2ea81bc 100755 --- a/utils/general.py +++ b/utils/general.py @@ -161,9 +161,15 @@ def emojis(str=''): return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str -def file_size(file): - # Return file size in MB - return Path(file).stat().st_size / 1e6 +def file_size(path): + # Return file/dir size (MB) + path = Path(path) + if path.is_file(): + return path.stat().st_size / 1E6 + elif path.is_dir(): + return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6 + else: + return 0.0 def check_online():