Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

export.py 15KB

il y a 3 ans
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>
il y a 3 ans
il y a 4 ans
il y a 4 ans
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337
  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. $ # Edit yolov5s_web_model/model.json to sort Identity* in ascending order
  15. $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
  16. $ npm install
  17. $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
  18. $ npm start
  19. """
  20. import argparse
  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. from models.common import Conv
  33. from models.experimental import attempt_load
  34. from models.yolo import Detect
  35. from utils.activations import SiLU
  36. from utils.datasets import LoadImages
  37. from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, print_args, \
  38. set_logging, url2file
  39. from utils.torch_utils import select_device
  40. def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
  41. # YOLOv5 TorchScript model export
  42. try:
  43. print(f'\n{prefix} starting export with torch {torch.__version__}...')
  44. f = file.with_suffix('.torchscript.pt')
  45. ts = torch.jit.trace(model, im, strict=False)
  46. (optimize_for_mobile(ts) if optimize else ts).save(f)
  47. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  48. except Exception as e:
  49. print(f'{prefix} export failure: {e}')
  50. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  51. # YOLOv5 ONNX export
  52. try:
  53. check_requirements(('onnx',))
  54. import onnx
  55. print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  56. f = file.with_suffix('.onnx')
  57. torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
  58. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  59. do_constant_folding=not train,
  60. input_names=['images'],
  61. output_names=['output'],
  62. dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
  63. 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
  64. } if dynamic else None)
  65. # Checks
  66. model_onnx = onnx.load(f) # load onnx model
  67. onnx.checker.check_model(model_onnx) # check onnx model
  68. # print(onnx.helper.printable_graph(model_onnx.graph)) # print
  69. # Simplify
  70. if simplify:
  71. try:
  72. check_requirements(('onnx-simplifier',))
  73. import onnxsim
  74. print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  75. model_onnx, check = onnxsim.simplify(
  76. model_onnx,
  77. dynamic_input_shape=dynamic,
  78. input_shapes={'images': list(im.shape)} if dynamic else None)
  79. assert check, 'assert check failed'
  80. onnx.save(model_onnx, f)
  81. except Exception as e:
  82. print(f'{prefix} simplifier failure: {e}')
  83. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  84. print(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
  85. except Exception as e:
  86. print(f'{prefix} export failure: {e}')
  87. def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
  88. # YOLOv5 CoreML export
  89. ct_model = None
  90. try:
  91. check_requirements(('coremltools',))
  92. import coremltools as ct
  93. print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  94. f = file.with_suffix('.mlmodel')
  95. model.train() # CoreML exports should be placed in model.train() mode
  96. ts = torch.jit.trace(model, im, strict=False) # TorchScript model
  97. ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255.0, bias=[0, 0, 0])])
  98. ct_model.save(f)
  99. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  100. except Exception as e:
  101. print(f'\n{prefix} export failure: {e}')
  102. return ct_model
  103. def export_saved_model(model, im, file, dynamic,
  104. tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
  105. conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
  106. # YOLOv5 TensorFlow saved_model export
  107. keras_model = None
  108. try:
  109. import tensorflow as tf
  110. from tensorflow import keras
  111. from models.tf import TFModel, TFDetect
  112. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  113. f = str(file).replace('.pt', '_saved_model')
  114. batch_size, ch, *imgsz = list(im.shape) # BCHW
  115. tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  116. im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
  117. y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  118. inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
  119. outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  120. keras_model = keras.Model(inputs=inputs, outputs=outputs)
  121. keras_model.trainable = False
  122. keras_model.summary()
  123. keras_model.save(f, save_format='tf')
  124. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  125. except Exception as e:
  126. print(f'\n{prefix} export failure: {e}')
  127. return keras_model
  128. def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
  129. # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  130. try:
  131. import tensorflow as tf
  132. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  133. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  134. f = file.with_suffix('.pb')
  135. m = tf.function(lambda x: keras_model(x)) # full model
  136. m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  137. frozen_func = convert_variables_to_constants_v2(m)
  138. frozen_func.graph.as_graph_def()
  139. tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  140. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  141. except Exception as e:
  142. print(f'\n{prefix} export failure: {e}')
  143. def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
  144. # YOLOv5 TensorFlow Lite export
  145. try:
  146. import tensorflow as tf
  147. from models.tf import representative_dataset_gen
  148. print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  149. batch_size, ch, *imgsz = list(im.shape) # BCHW
  150. f = str(file).replace('.pt', '-fp16.tflite')
  151. converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  152. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  153. converter.target_spec.supported_types = [tf.float16]
  154. converter.optimizations = [tf.lite.Optimize.DEFAULT]
  155. if int8:
  156. dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
  157. converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
  158. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  159. converter.target_spec.supported_types = []
  160. converter.inference_input_type = tf.uint8 # or tf.int8
  161. converter.inference_output_type = tf.uint8 # or tf.int8
  162. converter.experimental_new_quantizer = False
  163. f = str(file).replace('.pt', '-int8.tflite')
  164. tflite_model = converter.convert()
  165. open(f, "wb").write(tflite_model)
  166. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  167. except Exception as e:
  168. print(f'\n{prefix} export failure: {e}')
  169. def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
  170. # YOLOv5 TensorFlow.js export
  171. try:
  172. check_requirements(('tensorflowjs',))
  173. import tensorflowjs as tfjs
  174. print(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  175. f = str(file).replace('.pt', '_web_model') # js dir
  176. f_pb = file.with_suffix('.pb') # *.pb path
  177. cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \
  178. f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
  179. subprocess.run(cmd, shell=True)
  180. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  181. except Exception as e:
  182. print(f'\n{prefix} export failure: {e}')
  183. @torch.no_grad()
  184. def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
  185. weights=ROOT / 'yolov5s.pt', # weights path
  186. imgsz=(640, 640), # image (height, width)
  187. batch_size=1, # batch size
  188. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  189. include=('torchscript', 'onnx', 'coreml'), # include formats
  190. half=False, # FP16 half-precision export
  191. inplace=False, # set YOLOv5 Detect() inplace=True
  192. train=False, # model.train() mode
  193. optimize=False, # TorchScript: optimize for mobile
  194. int8=False, # CoreML/TF INT8 quantization
  195. dynamic=False, # ONNX/TF: dynamic axes
  196. simplify=False, # ONNX: simplify model
  197. opset=12, # ONNX: opset version
  198. ):
  199. t = time.time()
  200. include = [x.lower() for x in include]
  201. tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
  202. imgsz *= 2 if len(imgsz) == 1 else 1 # expand
  203. file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
  204. # Load PyTorch model
  205. device = select_device(device)
  206. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  207. model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
  208. nc, names = model.nc, model.names # number of classes, class names
  209. # Input
  210. gs = int(max(model.stride)) # grid size (max stride)
  211. imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
  212. im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
  213. # Update model
  214. if half:
  215. im, model = im.half(), model.half() # to FP16
  216. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  217. for k, m in model.named_modules():
  218. if isinstance(m, Conv): # assign export-friendly activations
  219. if isinstance(m.act, nn.SiLU):
  220. m.act = SiLU()
  221. elif isinstance(m, Detect):
  222. m.inplace = inplace
  223. m.onnx_dynamic = dynamic
  224. # m.forward = m.forward_export # assign forward (optional)
  225. for _ in range(2):
  226. y = model(im) # dry runs
  227. print(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
  228. # Exports
  229. if 'torchscript' in include:
  230. export_torchscript(model, im, file, optimize)
  231. if 'onnx' in include:
  232. export_onnx(model, im, file, opset, train, dynamic, simplify)
  233. if 'coreml' in include:
  234. export_coreml(model, im, file)
  235. # TensorFlow Exports
  236. if any(tf_exports):
  237. pb, tflite, tfjs = tf_exports[1:]
  238. assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
  239. model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs) # keras model
  240. if pb or tfjs: # pb prerequisite to tfjs
  241. export_pb(model, im, file)
  242. if tflite:
  243. export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
  244. if tfjs:
  245. export_tfjs(model, im, file)
  246. # Finish
  247. print(f'\nExport complete ({time.time() - t:.2f}s)'
  248. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  249. f'\nVisualize with https://netron.app')
  250. def parse_opt():
  251. parser = argparse.ArgumentParser()
  252. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  253. parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
  254. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  255. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  256. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  257. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  258. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  259. parser.add_argument('--train', action='store_true', help='model.train() mode')
  260. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  261. parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  262. parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
  263. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  264. parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version')
  265. parser.add_argument('--include', nargs='+',
  266. default=['torchscript', 'onnx'],
  267. help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)')
  268. opt = parser.parse_args()
  269. print_args(FILE.stem, opt)
  270. return opt
  271. def main(opt):
  272. set_logging()
  273. run(**vars(opt))
  274. if __name__ == "__main__":
  275. opt = parse_opt()
  276. main(opt)