You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

Merge `develop` branch into `master` (#3518) * update ci-testing.yml (#3322) * update ci-testing.yml * update greetings.yml * bring back os matrix * update ci-testing.yml (#3322) * update ci-testing.yml * update greetings.yml * bring back os matrix * Enable direct `--weights URL` definition (#3373) * Enable direct `--weights URL` definition @KalenMike this PR will enable direct --weights URL definition. Example use case: ``` python train.py --weights https://storage.googleapis.com/bucket/dir/model.pt ``` * cleanup * bug fixes * weights = attempt_download(weights) * Update experimental.py * Update hubconf.py * return bug fix * comment mirror * min_bytes * Update tutorial.ipynb (#3368) add Open in Kaggle badge * `cv2.imread(img, -1)` for IMREAD_UNCHANGED (#3379) * Update datasets.py * comment Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * COCO evolution fix (#3388) * COCO evolution fix * cleanup * update print * print fix * Create `is_pip()` function (#3391) Returns `True` if file is part of pip package. Useful for contextual behavior modification. ```python def is_pip(): # Is file in a pip package? return 'site-packages' in Path(__file__).absolute().parts ``` * Revert "`cv2.imread(img, -1)` for IMREAD_UNCHANGED (#3379)" (#3395) This reverts commit 21a9607e00f1365b21d8c4bd81bdbf5fc0efea24. * Update FLOPs description (#3422) * Update README.md * Changing FLOPS to FLOPs. Co-authored-by: BuildTools <unconfigured@null.spigotmc.org> * Parse URL authentication (#3424) * Parse URL authentication * urllib.parse.unquote() * improved error handling * improved error handling * remove %3F * update check_file() * Add FLOPs title to table (#3453) * Suppress jit trace warning + graph once (#3454) * Suppress jit trace warning + graph once Suppress harmless jit trace warning on TensorBoard add_graph call. Also fix multiple add_graph() calls bug, now only on batch 0. * Update train.py * Update MixUp augmentation `alpha=beta=32.0` (#3455) Per VOC empirical results https://github.com/ultralytics/yolov5/issues/3380#issuecomment-853001307 by @developer0hye * Add `timeout()` class (#3460) * Add `timeout()` class * rearrange order * Faster HSV augmentation (#3462) remove datatype conversion process that can be skipped * Add `check_git_status()` 5 second timeout (#3464) * Add check_git_status() 5 second timeout This should prevent the SSH Git bug that we were discussing @KalenMike * cleanup * replace timeout with check_output built-in timeout * Improved `check_requirements()` offline-handling (#3466) Improve robustness of `check_requirements()` function to offline environments (do not attempt pip installs when offline). * Add `output_names` argument for ONNX export with dynamic axes (#3456) * Add output names & dynamic axes for onnx export Add output_names and dynamic_axes names for all outputs in torch.onnx.export. The first four outputs of the model will have names output0, output1, output2, output3 * use first output only + cleanup Co-authored-by: Samridha Shrestha <samridha.shrestha@g42.ai> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Revert FP16 `test.py` and `detect.py` inference to FP32 default (#3423) * fixed inference bug ,while use half precision * replace --use-half with --half * replace space and PEP8 in detect.py * PEP8 detect.py * update --half help comment * Update test.py * revert space Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Add additional links/resources to stale.yml message (#3467) * Update stale.yml * cleanup * Update stale.yml * reformat * Update stale.yml HUB URL (#3468) * Stale `github.actor` bug fix (#3483) * Explicit `model.eval()` call `if opt.train=False` (#3475) * call model.eval() when opt.train is False call model.eval() when opt.train is False * single-line if statement * cleanup Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> * check_requirements() exclude `opencv-python` (#3495) Fix for 3rd party or contrib versions of installed OpenCV as in https://github.com/ultralytics/yolov5/issues/3494. * Earlier `assert` for cpu and half option (#3508) * early assert for cpu and half option early assert for cpu and half option * Modified comment Modified comment * Update tutorial.ipynb (#3510) * Reduce test.py results spacing (#3511) * Update README.md (#3512) * Update README.md Minor modifications * 850 width * Update greetings.yml revert greeting change as PRs will now merge to master. Co-authored-by: Piotr Skalski <SkalskiP@users.noreply.github.com> Co-authored-by: SkalskiP <piotr.skalski92@gmail.com> Co-authored-by: Peretz Cohen <pizzaz93@users.noreply.github.com> Co-authored-by: tudoulei <34886368+tudoulei@users.noreply.github.com> Co-authored-by: chocosaj <chocosaj@users.noreply.github.com> Co-authored-by: BuildTools <unconfigured@null.spigotmc.org> Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com> Co-authored-by: Sam_S <SamSamhuns@users.noreply.github.com> Co-authored-by: Samridha Shrestha <samridha.shrestha@g42.ai> Co-authored-by: edificewang <609552430@qq.com>
3 anni fa
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364
  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Export a YOLOv5 PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats
  4. TensorFlow exports authored by https://github.com/zldrobit
  5. Usage:
  6. $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
  7. Inference:
  8. $ python path/to/detect.py --weights yolov5s.pt
  9. yolov5s.onnx (must export with --dynamic)
  10. yolov5s_saved_model
  11. yolov5s.pb
  12. yolov5s.tflite
  13. TensorFlow.js:
  14. $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
  15. $ npm install
  16. $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
  17. $ npm start
  18. """
  19. import argparse
  20. import os
  21. import subprocess
  22. import sys
  23. import time
  24. from pathlib import Path
  25. import torch
  26. import torch.nn as nn
  27. from torch.utils.mobile_optimizer import optimize_for_mobile
  28. FILE = Path(__file__).resolve()
  29. ROOT = FILE.parents[0] # YOLOv5 root directory
  30. if str(ROOT) not in sys.path:
  31. sys.path.append(str(ROOT)) # add ROOT to PATH
  32. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  33. from models.common import Conv
  34. from models.experimental import attempt_load
  35. from models.yolo import Detect
  36. from utils.activations import SiLU
  37. from utils.datasets import LoadImages
  38. from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, print_args, \
  39. set_logging, url2file
  40. from utils.torch_utils import select_device
  41. def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
  42. # YOLOv5 TorchScript model export
  43. try:
  44. print(f'\n{prefix} starting export with torch {torch.__version__}...')
  45. f = file.with_suffix('.torchscript.pt')
  46. ts = torch.jit.trace(model, im, strict=False)
  47. (optimize_for_mobile(ts) if optimize else ts).save(f)
  48. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  49. except Exception as e:
  50. print(f'{prefix} export failure: {e}')
  51. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  52. # YOLOv5 ONNX export
  53. try:
  54. check_requirements(('onnx',))
  55. import onnx
  56. print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  57. f = file.with_suffix('.onnx')
  58. torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
  59. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  60. do_constant_folding=not train,
  61. input_names=['images'],
  62. output_names=['output'],
  63. dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
  64. 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
  65. } if dynamic else None)
  66. # Checks
  67. model_onnx = onnx.load(f) # load onnx model
  68. onnx.checker.check_model(model_onnx) # check onnx model
  69. # print(onnx.helper.printable_graph(model_onnx.graph)) # print
  70. # Simplify
  71. if simplify:
  72. try:
  73. check_requirements(('onnx-simplifier',))
  74. import onnxsim
  75. print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  76. model_onnx, check = onnxsim.simplify(
  77. model_onnx,
  78. dynamic_input_shape=dynamic,
  79. input_shapes={'images': list(im.shape)} if dynamic else None)
  80. assert check, 'assert check failed'
  81. onnx.save(model_onnx, f)
  82. except Exception as e:
  83. print(f'{prefix} simplifier failure: {e}')
  84. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  85. print(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
  86. except Exception as e:
  87. print(f'{prefix} export failure: {e}')
  88. def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
  89. # YOLOv5 CoreML export
  90. ct_model = None
  91. try:
  92. check_requirements(('coremltools',))
  93. import coremltools as ct
  94. print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  95. f = file.with_suffix('.mlmodel')
  96. model.train() # CoreML exports should be placed in model.train() mode
  97. ts = torch.jit.trace(model, im, strict=False) # TorchScript model
  98. ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255.0, bias=[0, 0, 0])])
  99. ct_model.save(f)
  100. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  101. except Exception as e:
  102. print(f'\n{prefix} export failure: {e}')
  103. return ct_model
  104. def export_saved_model(model, im, file, dynamic,
  105. tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
  106. conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
  107. # YOLOv5 TensorFlow saved_model export
  108. keras_model = None
  109. try:
  110. import tensorflow as tf
  111. from tensorflow import keras
  112. from models.tf import TFModel, TFDetect
  113. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  114. f = str(file).replace('.pt', '_saved_model')
  115. batch_size, ch, *imgsz = list(im.shape) # BCHW
  116. tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  117. im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
  118. y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  119. inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
  120. outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  121. keras_model = keras.Model(inputs=inputs, outputs=outputs)
  122. keras_model.trainable = False
  123. keras_model.summary()
  124. keras_model.save(f, save_format='tf')
  125. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  126. except Exception as e:
  127. print(f'\n{prefix} export failure: {e}')
  128. return keras_model
  129. def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
  130. # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  131. try:
  132. import tensorflow as tf
  133. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  134. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  135. f = file.with_suffix('.pb')
  136. m = tf.function(lambda x: keras_model(x)) # full model
  137. m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  138. frozen_func = convert_variables_to_constants_v2(m)
  139. frozen_func.graph.as_graph_def()
  140. tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  141. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  142. except Exception as e:
  143. print(f'\n{prefix} export failure: {e}')
  144. def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
  145. # YOLOv5 TensorFlow Lite export
  146. try:
  147. import tensorflow as tf
  148. from models.tf import representative_dataset_gen
  149. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  150. batch_size, ch, *imgsz = list(im.shape) # BCHW
  151. f = str(file).replace('.pt', '-fp16.tflite')
  152. converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  153. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  154. converter.target_spec.supported_types = [tf.float16]
  155. converter.optimizations = [tf.lite.Optimize.DEFAULT]
  156. if int8:
  157. dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
  158. converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
  159. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  160. converter.target_spec.supported_types = []
  161. converter.inference_input_type = tf.uint8 # or tf.int8
  162. converter.inference_output_type = tf.uint8 # or tf.int8
  163. converter.experimental_new_quantizer = False
  164. f = str(file).replace('.pt', '-int8.tflite')
  165. tflite_model = converter.convert()
  166. open(f, "wb").write(tflite_model)
  167. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  168. except Exception as e:
  169. print(f'\n{prefix} export failure: {e}')
  170. def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
  171. # YOLOv5 TensorFlow.js export
  172. try:
  173. check_requirements(('tensorflowjs',))
  174. import re
  175. import tensorflowjs as tfjs
  176. print(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  177. f = str(file).replace('.pt', '_web_model') # js dir
  178. f_pb = file.with_suffix('.pb') # *.pb path
  179. f_json = f + '/model.json' # *.json path
  180. cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \
  181. f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
  182. subprocess.run(cmd, shell=True)
  183. json = open(f_json).read()
  184. with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
  185. subst = re.sub(
  186. r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
  187. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  188. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  189. r'"Identity.?.?": {"name": "Identity.?.?"}}}',
  190. r'{"outputs": {"Identity": {"name": "Identity"}, '
  191. r'"Identity_1": {"name": "Identity_1"}, '
  192. r'"Identity_2": {"name": "Identity_2"}, '
  193. r'"Identity_3": {"name": "Identity_3"}}}',
  194. json)
  195. j.write(subst)
  196. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  197. except Exception as e:
  198. print(f'\n{prefix} export failure: {e}')
  199. @torch.no_grad()
  200. def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
  201. weights=ROOT / 'yolov5s.pt', # weights path
  202. imgsz=(640, 640), # image (height, width)
  203. batch_size=1, # batch size
  204. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  205. include=('torchscript', 'onnx', 'coreml'), # include formats
  206. half=False, # FP16 half-precision export
  207. inplace=False, # set YOLOv5 Detect() inplace=True
  208. train=False, # model.train() mode
  209. optimize=False, # TorchScript: optimize for mobile
  210. int8=False, # CoreML/TF INT8 quantization
  211. dynamic=False, # ONNX/TF: dynamic axes
  212. simplify=False, # ONNX: simplify model
  213. opset=12, # ONNX: opset version
  214. topk_per_class=100, # TF.js NMS: topk per class to keep
  215. topk_all=100, # TF.js NMS: topk for all classes to keep
  216. iou_thres=0.45, # TF.js NMS: IoU threshold
  217. conf_thres=0.25 # TF.js NMS: confidence threshold
  218. ):
  219. t = time.time()
  220. include = [x.lower() for x in include]
  221. tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
  222. imgsz *= 2 if len(imgsz) == 1 else 1 # expand
  223. file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
  224. # Load PyTorch model
  225. device = select_device(device)
  226. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  227. model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
  228. nc, names = model.nc, model.names # number of classes, class names
  229. # Input
  230. gs = int(max(model.stride)) # grid size (max stride)
  231. imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
  232. im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
  233. # Update model
  234. if half:
  235. im, model = im.half(), model.half() # to FP16
  236. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  237. for k, m in model.named_modules():
  238. if isinstance(m, Conv): # assign export-friendly activations
  239. if isinstance(m.act, nn.SiLU):
  240. m.act = SiLU()
  241. elif isinstance(m, Detect):
  242. m.inplace = inplace
  243. m.onnx_dynamic = dynamic
  244. # m.forward = m.forward_export # assign forward (optional)
  245. for _ in range(2):
  246. y = model(im) # dry runs
  247. print(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
  248. # Exports
  249. if 'torchscript' in include:
  250. export_torchscript(model, im, file, optimize)
  251. if 'onnx' in include:
  252. export_onnx(model, im, file, opset, train, dynamic, simplify)
  253. if 'coreml' in include:
  254. export_coreml(model, im, file)
  255. # TensorFlow Exports
  256. if any(tf_exports):
  257. pb, tflite, tfjs = tf_exports[1:]
  258. assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
  259. model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs,
  260. topk_per_class=topk_per_class, topk_all=topk_all, conf_thres=conf_thres,
  261. iou_thres=iou_thres) # keras model
  262. if pb or tfjs: # pb prerequisite to tfjs
  263. export_pb(model, im, file)
  264. if tflite:
  265. export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
  266. if tfjs:
  267. export_tfjs(model, im, file)
  268. # Finish
  269. print(f'\nExport complete ({time.time() - t:.2f}s)'
  270. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  271. f'\nVisualize with https://netron.app')
  272. def parse_opt():
  273. parser = argparse.ArgumentParser()
  274. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  275. parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
  276. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  277. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  278. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  279. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  280. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  281. parser.add_argument('--train', action='store_true', help='model.train() mode')
  282. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  283. parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  284. parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
  285. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  286. parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version')
  287. parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
  288. parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
  289. parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
  290. parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
  291. parser.add_argument('--include', nargs='+',
  292. default=['torchscript', 'onnx'],
  293. help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)')
  294. opt = parser.parse_args()
  295. print_args(FILE.stem, opt)
  296. return opt
  297. def main(opt):
  298. set_logging()
  299. run(**vars(opt))
  300. if __name__ == "__main__":
  301. opt = parse_opt()
  302. main(opt)