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.

608 lines
29KB

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