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.

589 lines
28KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
  4. Format | `export.py --include` | Model
  5. --- | --- | ---
  6. PyTorch | - | yolov5s.pt
  7. TorchScript | `torchscript` | yolov5s.torchscript
  8. ONNX | `onnx` | yolov5s.onnx
  9. OpenVINO | `openvino` | yolov5s_openvino_model/
  10. TensorRT | `engine` | yolov5s.engine
  11. CoreML | `coreml` | yolov5s.mlmodel
  12. TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/
  13. TensorFlow GraphDef | `pb` | yolov5s.pb
  14. TensorFlow Lite | `tflite` | yolov5s.tflite
  15. TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
  16. TensorFlow.js | `tfjs` | yolov5s_web_model/
  17. Requirements:
  18. $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU
  19. $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
  20. Usage:
  21. $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
  22. Inference:
  23. $ python path/to/detect.py --weights yolov5s.pt # PyTorch
  24. yolov5s.torchscript # TorchScript
  25. yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
  26. yolov5s.xml # OpenVINO
  27. yolov5s.engine # TensorRT
  28. yolov5s.mlmodel # CoreML (macOS-only)
  29. yolov5s_saved_model # TensorFlow SavedModel
  30. yolov5s.pb # TensorFlow GraphDef
  31. yolov5s.tflite # TensorFlow Lite
  32. yolov5s_edgetpu.tflite # TensorFlow Edge TPU
  33. TensorFlow.js:
  34. $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
  35. $ npm install
  36. $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
  37. $ npm start
  38. """
  39. import argparse
  40. import json
  41. import os
  42. import platform
  43. import subprocess
  44. import sys
  45. import time
  46. import warnings
  47. from pathlib import Path
  48. import pandas as pd
  49. import torch
  50. from torch.utils.mobile_optimizer import optimize_for_mobile
  51. FILE = Path(__file__).resolve()
  52. ROOT = FILE.parents[0] # YOLOv5 root directory
  53. if str(ROOT) not in sys.path:
  54. sys.path.append(str(ROOT)) # add ROOT to PATH
  55. if platform.system() != 'Windows':
  56. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  57. from models.experimental import attempt_load
  58. from models.yolo import Detect
  59. from utils.datasets import LoadImages
  60. from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr,
  61. file_size, print_args, url2file)
  62. from utils.torch_utils import select_device
  63. def export_formats():
  64. # YOLOv5 export formats
  65. x = [
  66. ['PyTorch', '-', '.pt', True],
  67. ['TorchScript', 'torchscript', '.torchscript', True],
  68. ['ONNX', 'onnx', '.onnx', True],
  69. ['OpenVINO', 'openvino', '_openvino_model', False],
  70. ['TensorRT', 'engine', '.engine', True],
  71. ['CoreML', 'coreml', '.mlmodel', False],
  72. ['TensorFlow SavedModel', 'saved_model', '_saved_model', True],
  73. ['TensorFlow GraphDef', 'pb', '.pb', True],
  74. ['TensorFlow Lite', 'tflite', '.tflite', False],
  75. ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False],
  76. ['TensorFlow.js', 'tfjs', '_web_model', False],]
  77. return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'GPU'])
  78. def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
  79. # YOLOv5 TorchScript model export
  80. try:
  81. LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
  82. f = file.with_suffix('.torchscript')
  83. ts = torch.jit.trace(model, im, strict=False)
  84. d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
  85. extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
  86. if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
  87. optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
  88. else:
  89. ts.save(str(f), _extra_files=extra_files)
  90. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  91. return f
  92. except Exception as e:
  93. LOGGER.info(f'{prefix} export failure: {e}')
  94. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  95. # YOLOv5 ONNX export
  96. try:
  97. check_requirements(('onnx',))
  98. import onnx
  99. LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  100. f = file.with_suffix('.onnx')
  101. torch.onnx.export(
  102. model,
  103. im,
  104. f,
  105. verbose=False,
  106. opset_version=opset,
  107. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  108. do_constant_folding=not train,
  109. input_names=['images'],
  110. output_names=['output'],
  111. dynamic_axes={
  112. 'images': {
  113. 0: 'batch',
  114. 2: 'height',
  115. 3: 'width'}, # shape(1,3,640,640)
  116. 'output': {
  117. 0: 'batch',
  118. 1: 'anchors'} # shape(1,25200,85)
  119. } if dynamic else None)
  120. # Checks
  121. model_onnx = onnx.load(f) # load onnx model
  122. onnx.checker.check_model(model_onnx) # check onnx model
  123. # Metadata
  124. d = {'stride': int(max(model.stride)), 'names': model.names}
  125. for k, v in d.items():
  126. meta = model_onnx.metadata_props.add()
  127. meta.key, meta.value = k, str(v)
  128. onnx.save(model_onnx, f)
  129. # Simplify
  130. if simplify:
  131. try:
  132. check_requirements(('onnx-simplifier',))
  133. import onnxsim
  134. LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  135. model_onnx, check = onnxsim.simplify(model_onnx,
  136. dynamic_input_shape=dynamic,
  137. input_shapes={'images': list(im.shape)} if dynamic else None)
  138. assert check, 'assert check failed'
  139. onnx.save(model_onnx, f)
  140. except Exception as e:
  141. LOGGER.info(f'{prefix} simplifier failure: {e}')
  142. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  143. return f
  144. except Exception as e:
  145. LOGGER.info(f'{prefix} export failure: {e}')
  146. def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
  147. # YOLOv5 OpenVINO export
  148. try:
  149. check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
  150. import openvino.inference_engine as ie
  151. LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
  152. f = str(file).replace('.pt', '_openvino_model' + os.sep)
  153. cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}"
  154. subprocess.check_output(cmd, shell=True)
  155. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  156. return f
  157. except Exception as e:
  158. LOGGER.info(f'\n{prefix} export failure: {e}')
  159. def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
  160. # YOLOv5 CoreML export
  161. try:
  162. check_requirements(('coremltools',))
  163. import coremltools as ct
  164. LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  165. f = file.with_suffix('.mlmodel')
  166. ts = torch.jit.trace(model, im, strict=False) # TorchScript model
  167. ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
  168. ct_model.save(f)
  169. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  170. return ct_model, f
  171. except Exception as e:
  172. LOGGER.info(f'\n{prefix} export failure: {e}')
  173. return None, None
  174. def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
  175. # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
  176. try:
  177. check_requirements(('tensorrt',))
  178. import tensorrt as trt
  179. if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
  180. grid = model.model[-1].anchor_grid
  181. model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
  182. export_onnx(model, im, file, 12, train, False, simplify) # opset 12
  183. model.model[-1].anchor_grid = grid
  184. else: # TensorRT >= 8
  185. check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0
  186. export_onnx(model, im, file, 13, train, False, simplify) # opset 13
  187. onnx = file.with_suffix('.onnx')
  188. LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
  189. assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
  190. assert onnx.exists(), f'failed to export ONNX file: {onnx}'
  191. f = file.with_suffix('.engine') # TensorRT engine file
  192. logger = trt.Logger(trt.Logger.INFO)
  193. if verbose:
  194. logger.min_severity = trt.Logger.Severity.VERBOSE
  195. builder = trt.Builder(logger)
  196. config = builder.create_builder_config()
  197. config.max_workspace_size = workspace * 1 << 30
  198. # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice
  199. flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
  200. network = builder.create_network(flag)
  201. parser = trt.OnnxParser(network, logger)
  202. if not parser.parse_from_file(str(onnx)):
  203. raise RuntimeError(f'failed to load ONNX file: {onnx}')
  204. inputs = [network.get_input(i) for i in range(network.num_inputs)]
  205. outputs = [network.get_output(i) for i in range(network.num_outputs)]
  206. LOGGER.info(f'{prefix} Network Description:')
  207. for inp in inputs:
  208. LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
  209. for out in outputs:
  210. LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
  211. LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 else 32} engine in {f}')
  212. if builder.platform_has_fast_fp16:
  213. config.set_flag(trt.BuilderFlag.FP16)
  214. with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
  215. t.write(engine.serialize())
  216. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  217. return f
  218. except Exception as e:
  219. LOGGER.info(f'\n{prefix} export failure: {e}')
  220. def export_saved_model(model,
  221. im,
  222. file,
  223. dynamic,
  224. tf_nms=False,
  225. agnostic_nms=False,
  226. topk_per_class=100,
  227. topk_all=100,
  228. iou_thres=0.45,
  229. conf_thres=0.25,
  230. keras=False,
  231. prefix=colorstr('TensorFlow SavedModel:')):
  232. # YOLOv5 TensorFlow SavedModel export
  233. try:
  234. import tensorflow as tf
  235. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  236. from models.tf import TFDetect, TFModel
  237. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  238. f = str(file).replace('.pt', '_saved_model')
  239. batch_size, ch, *imgsz = list(im.shape) # BCHW
  240. tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  241. im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow
  242. _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  243. inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
  244. outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  245. keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
  246. keras_model.trainable = False
  247. keras_model.summary()
  248. if keras:
  249. keras_model.save(f, save_format='tf')
  250. else:
  251. spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
  252. m = tf.function(lambda x: keras_model(x)) # full model
  253. m = m.get_concrete_function(spec)
  254. frozen_func = convert_variables_to_constants_v2(m)
  255. tfm = tf.Module()
  256. tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x)[0], [spec])
  257. tfm.__call__(im)
  258. tf.saved_model.save(tfm,
  259. f,
  260. options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
  261. if check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions())
  262. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  263. return keras_model, f
  264. except Exception as e:
  265. LOGGER.info(f'\n{prefix} export failure: {e}')
  266. return None, None
  267. def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
  268. # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  269. try:
  270. import tensorflow as tf
  271. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  272. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  273. f = file.with_suffix('.pb')
  274. m = tf.function(lambda x: keras_model(x)) # full model
  275. m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  276. frozen_func = convert_variables_to_constants_v2(m)
  277. frozen_func.graph.as_graph_def()
  278. tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  279. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  280. return f
  281. except Exception as e:
  282. LOGGER.info(f'\n{prefix} export failure: {e}')
  283. def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
  284. # YOLOv5 TensorFlow Lite export
  285. try:
  286. import tensorflow as tf
  287. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  288. batch_size, ch, *imgsz = list(im.shape) # BCHW
  289. f = str(file).replace('.pt', '-fp16.tflite')
  290. converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  291. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  292. converter.target_spec.supported_types = [tf.float16]
  293. converter.optimizations = [tf.lite.Optimize.DEFAULT]
  294. if int8:
  295. from models.tf import representative_dataset_gen
  296. dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
  297. converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
  298. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  299. converter.target_spec.supported_types = []
  300. converter.inference_input_type = tf.uint8 # or tf.int8
  301. converter.inference_output_type = tf.uint8 # or tf.int8
  302. converter.experimental_new_quantizer = True
  303. f = str(file).replace('.pt', '-int8.tflite')
  304. if nms or agnostic_nms:
  305. converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
  306. tflite_model = converter.convert()
  307. open(f, "wb").write(tflite_model)
  308. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  309. return f
  310. except Exception as e:
  311. LOGGER.info(f'\n{prefix} export failure: {e}')
  312. def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
  313. # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
  314. try:
  315. cmd = 'edgetpu_compiler --version'
  316. help_url = 'https://coral.ai/docs/edgetpu/compiler/'
  317. assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'
  318. if subprocess.run(cmd + ' >/dev/null', shell=True).returncode != 0:
  319. LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
  320. sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system
  321. for c in (
  322. 'curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
  323. 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
  324. 'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'):
  325. subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)
  326. ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
  327. LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
  328. f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model
  329. f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model
  330. cmd = f"edgetpu_compiler -s {f_tfl}"
  331. subprocess.run(cmd, shell=True, check=True)
  332. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  333. return f
  334. except Exception as e:
  335. LOGGER.info(f'\n{prefix} export failure: {e}')
  336. def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
  337. # YOLOv5 TensorFlow.js export
  338. try:
  339. check_requirements(('tensorflowjs',))
  340. import re
  341. import tensorflowjs as tfjs
  342. LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  343. f = str(file).replace('.pt', '_web_model') # js dir
  344. f_pb = file.with_suffix('.pb') # *.pb path
  345. f_json = f + '/model.json' # *.json path
  346. cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \
  347. f'--output_node_names="Identity,Identity_1,Identity_2,Identity_3" {f_pb} {f}'
  348. subprocess.run(cmd, shell=True)
  349. with open(f_json) as j:
  350. json = j.read()
  351. with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
  352. subst = re.sub(
  353. r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
  354. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  355. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  356. r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
  357. r'"Identity_1": {"name": "Identity_1"}, '
  358. r'"Identity_2": {"name": "Identity_2"}, '
  359. r'"Identity_3": {"name": "Identity_3"}}}', json)
  360. j.write(subst)
  361. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  362. return f
  363. except Exception as e:
  364. LOGGER.info(f'\n{prefix} export failure: {e}')
  365. @torch.no_grad()
  366. def run(
  367. data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
  368. weights=ROOT / 'yolov5s.pt', # weights path
  369. imgsz=(640, 640), # image (height, width)
  370. batch_size=1, # batch size
  371. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  372. include=('torchscript', 'onnx'), # include formats
  373. half=False, # FP16 half-precision export
  374. inplace=False, # set YOLOv5 Detect() inplace=True
  375. train=False, # model.train() mode
  376. optimize=False, # TorchScript: optimize for mobile
  377. int8=False, # CoreML/TF INT8 quantization
  378. dynamic=False, # ONNX/TF: dynamic axes
  379. simplify=False, # ONNX: simplify model
  380. opset=12, # ONNX: opset version
  381. verbose=False, # TensorRT: verbose log
  382. workspace=4, # TensorRT: workspace size (GB)
  383. nms=False, # TF: add NMS to model
  384. agnostic_nms=False, # TF: add agnostic NMS to model
  385. topk_per_class=100, # TF.js NMS: topk per class to keep
  386. topk_all=100, # TF.js NMS: topk for all classes to keep
  387. iou_thres=0.45, # TF.js NMS: IoU threshold
  388. conf_thres=0.25, # TF.js NMS: confidence threshold
  389. ):
  390. t = time.time()
  391. include = [x.lower() for x in include] # to lowercase
  392. formats = tuple(export_formats()['Argument'][1:]) # --include arguments
  393. flags = [x in include for x in formats]
  394. assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {formats}'
  395. jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = flags # export booleans
  396. file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) # PyTorch weights
  397. # Load PyTorch model
  398. device = select_device(device)
  399. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  400. model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
  401. nc, names = model.nc, model.names # number of classes, class names
  402. # Checks
  403. imgsz *= 2 if len(imgsz) == 1 else 1 # expand
  404. opset = 12 if ('openvino' in include) else opset # OpenVINO requires opset <= 12
  405. assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
  406. # Input
  407. gs = int(max(model.stride)) # grid size (max stride)
  408. imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
  409. im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
  410. # Update model
  411. if half:
  412. im, model = im.half(), model.half() # to FP16
  413. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  414. for k, m in model.named_modules():
  415. if isinstance(m, Detect):
  416. m.inplace = inplace
  417. m.onnx_dynamic = dynamic
  418. m.export = True
  419. for _ in range(2):
  420. y = model(im) # dry runs
  421. shape = tuple(y[0].shape) # model output shape
  422. LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
  423. # Exports
  424. f = [''] * 10 # exported filenames
  425. warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
  426. if jit:
  427. f[0] = export_torchscript(model, im, file, optimize)
  428. if engine: # TensorRT required before ONNX
  429. f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose)
  430. if onnx or xml: # OpenVINO requires ONNX
  431. f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
  432. if xml: # OpenVINO
  433. f[3] = export_openvino(model, im, file)
  434. if coreml:
  435. _, f[4] = export_coreml(model, im, file)
  436. # TensorFlow Exports
  437. if any((saved_model, pb, tflite, edgetpu, tfjs)):
  438. if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
  439. check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
  440. assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
  441. model, f[5] = export_saved_model(model.cpu(),
  442. im,
  443. file,
  444. dynamic,
  445. tf_nms=nms or agnostic_nms or tfjs,
  446. agnostic_nms=agnostic_nms or tfjs,
  447. topk_per_class=topk_per_class,
  448. topk_all=topk_all,
  449. conf_thres=conf_thres,
  450. iou_thres=iou_thres) # keras model
  451. if pb or tfjs: # pb prerequisite to tfjs
  452. f[6] = export_pb(model, im, file)
  453. if tflite or edgetpu:
  454. f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms)
  455. if edgetpu:
  456. f[8] = export_edgetpu(model, im, file)
  457. if tfjs:
  458. f[9] = export_tfjs(model, im, file)
  459. # Finish
  460. f = [str(x) for x in f if x] # filter out '' and None
  461. if any(f):
  462. LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
  463. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  464. f"\nDetect: python detect.py --weights {f[-1]}"
  465. f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')"
  466. f"\nValidate: python val.py --weights {f[-1]}"
  467. f"\nVisualize: https://netron.app")
  468. return f # return list of exported files/dirs
  469. def parse_opt():
  470. parser = argparse.ArgumentParser()
  471. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  472. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
  473. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  474. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  475. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  476. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  477. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  478. parser.add_argument('--train', action='store_true', help='model.train() mode')
  479. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  480. parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  481. parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
  482. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  483. parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
  484. parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
  485. parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
  486. parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
  487. parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
  488. parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
  489. parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
  490. parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
  491. parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
  492. parser.add_argument('--include',
  493. nargs='+',
  494. default=['torchscript', 'onnx'],
  495. help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
  496. opt = parser.parse_args()
  497. print_args(vars(opt))
  498. return opt
  499. def main(opt):
  500. for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
  501. run(**vars(opt))
  502. if __name__ == "__main__":
  503. opt = parse_opt()
  504. main(opt)