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.

560 lignes
27KB

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