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.

441 lines
20KB

  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 | Example | `--include ...` argument
  5. --- | --- | ---
  6. PyTorch | yolov5s.pt | -
  7. TorchScript | yolov5s.torchscript | `torchscript`
  8. ONNX | yolov5s.onnx | `onnx`
  9. CoreML | yolov5s.mlmodel | `coreml`
  10. TensorFlow SavedModel | yolov5s_saved_model/ | `saved_model`
  11. TensorFlow GraphDef | yolov5s.pb | `pb`
  12. TensorFlow Lite | yolov5s.tflite | `tflite`
  13. TensorFlow.js | yolov5s_web_model/ | `tfjs`
  14. TensorRT | yolov5s.engine | `engine`
  15. Usage:
  16. $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
  17. Inference:
  18. $ python path/to/detect.py --weights yolov5s.pt
  19. yolov5s.torchscript
  20. yolov5s.onnx
  21. yolov5s.mlmodel (under development)
  22. yolov5s_saved_model
  23. yolov5s.pb
  24. yolov5s.tflite
  25. yolov5s.engine
  26. TensorFlow.js:
  27. $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
  28. $ npm install
  29. $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
  30. $ npm start
  31. """
  32. import argparse
  33. import json
  34. import os
  35. import subprocess
  36. import sys
  37. import time
  38. from pathlib import Path
  39. import torch
  40. import torch.nn as nn
  41. from torch.utils.mobile_optimizer import optimize_for_mobile
  42. FILE = Path(__file__).resolve()
  43. ROOT = FILE.parents[0] # YOLOv5 root directory
  44. if str(ROOT) not in sys.path:
  45. sys.path.append(str(ROOT)) # add ROOT to PATH
  46. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  47. from models.common import Conv
  48. from models.experimental import attempt_load
  49. from models.yolo import Detect
  50. from utils.activations import SiLU
  51. from utils.datasets import LoadImages
  52. from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, colorstr, file_size, print_args,
  53. url2file)
  54. from utils.torch_utils import select_device
  55. def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
  56. # YOLOv5 TorchScript model export
  57. try:
  58. LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
  59. f = file.with_suffix('.torchscript')
  60. ts = torch.jit.trace(model, im, strict=False)
  61. d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
  62. extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
  63. (optimize_for_mobile(ts) if optimize else ts).save(str(f), _extra_files=extra_files)
  64. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  65. except Exception as e:
  66. LOGGER.info(f'{prefix} export failure: {e}')
  67. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  68. # YOLOv5 ONNX export
  69. try:
  70. check_requirements(('onnx',))
  71. import onnx
  72. LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  73. f = file.with_suffix('.onnx')
  74. torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
  75. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  76. do_constant_folding=not train,
  77. input_names=['images'],
  78. output_names=['output'],
  79. dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
  80. 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
  81. } if dynamic else None)
  82. # Checks
  83. model_onnx = onnx.load(f) # load onnx model
  84. onnx.checker.check_model(model_onnx) # check onnx model
  85. # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph)) # print
  86. # Simplify
  87. if simplify:
  88. try:
  89. check_requirements(('onnx-simplifier',))
  90. import onnxsim
  91. LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  92. model_onnx, check = onnxsim.simplify(
  93. model_onnx,
  94. dynamic_input_shape=dynamic,
  95. input_shapes={'images': list(im.shape)} if dynamic else None)
  96. assert check, 'assert check failed'
  97. onnx.save(model_onnx, f)
  98. except Exception as e:
  99. LOGGER.info(f'{prefix} simplifier failure: {e}')
  100. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  101. LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
  102. except Exception as e:
  103. LOGGER.info(f'{prefix} export failure: {e}')
  104. def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
  105. # YOLOv5 CoreML export
  106. ct_model = None
  107. try:
  108. check_requirements(('coremltools',))
  109. import coremltools as ct
  110. LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  111. f = file.with_suffix('.mlmodel')
  112. model.train() # CoreML exports should be placed in model.train() mode
  113. ts = torch.jit.trace(model, im, strict=False) # TorchScript model
  114. ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
  115. ct_model.save(f)
  116. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  117. except Exception as e:
  118. LOGGER.info(f'\n{prefix} export failure: {e}')
  119. return ct_model
  120. def export_saved_model(model, im, file, dynamic,
  121. tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
  122. conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
  123. # YOLOv5 TensorFlow saved_model export
  124. keras_model = None
  125. try:
  126. import tensorflow as tf
  127. from tensorflow import keras
  128. from models.tf import TFDetect, TFModel
  129. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  130. f = str(file).replace('.pt', '_saved_model')
  131. batch_size, ch, *imgsz = list(im.shape) # BCHW
  132. tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  133. im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
  134. y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  135. inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
  136. outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  137. keras_model = keras.Model(inputs=inputs, outputs=outputs)
  138. keras_model.trainable = False
  139. keras_model.summary()
  140. keras_model.save(f, save_format='tf')
  141. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  142. except Exception as e:
  143. LOGGER.info(f'\n{prefix} export failure: {e}')
  144. return keras_model
  145. def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
  146. # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  147. try:
  148. import tensorflow as tf
  149. from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  150. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  151. f = file.with_suffix('.pb')
  152. m = tf.function(lambda x: keras_model(x)) # full model
  153. m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  154. frozen_func = convert_variables_to_constants_v2(m)
  155. frozen_func.graph.as_graph_def()
  156. tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  157. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  158. except Exception as e:
  159. LOGGER.info(f'\n{prefix} export failure: {e}')
  160. def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
  161. # YOLOv5 TensorFlow Lite export
  162. try:
  163. import tensorflow as tf
  164. from models.tf import representative_dataset_gen
  165. LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  166. batch_size, ch, *imgsz = list(im.shape) # BCHW
  167. f = str(file).replace('.pt', '-fp16.tflite')
  168. converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  169. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  170. converter.target_spec.supported_types = [tf.float16]
  171. converter.optimizations = [tf.lite.Optimize.DEFAULT]
  172. if int8:
  173. dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
  174. converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
  175. converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  176. converter.target_spec.supported_types = []
  177. converter.inference_input_type = tf.uint8 # or tf.int8
  178. converter.inference_output_type = tf.uint8 # or tf.int8
  179. converter.experimental_new_quantizer = False
  180. f = str(file).replace('.pt', '-int8.tflite')
  181. tflite_model = converter.convert()
  182. open(f, "wb").write(tflite_model)
  183. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  184. except Exception as e:
  185. LOGGER.info(f'\n{prefix} export failure: {e}')
  186. def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
  187. # YOLOv5 TensorFlow.js export
  188. try:
  189. check_requirements(('tensorflowjs',))
  190. import re
  191. import tensorflowjs as tfjs
  192. LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  193. f = str(file).replace('.pt', '_web_model') # js dir
  194. f_pb = file.with_suffix('.pb') # *.pb path
  195. f_json = f + '/model.json' # *.json path
  196. cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \
  197. f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
  198. subprocess.run(cmd, shell=True)
  199. json = open(f_json).read()
  200. with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
  201. subst = re.sub(
  202. r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
  203. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  204. r'"Identity.?.?": {"name": "Identity.?.?"}, '
  205. r'"Identity.?.?": {"name": "Identity.?.?"}}}',
  206. r'{"outputs": {"Identity": {"name": "Identity"}, '
  207. r'"Identity_1": {"name": "Identity_1"}, '
  208. r'"Identity_2": {"name": "Identity_2"}, '
  209. r'"Identity_3": {"name": "Identity_3"}}}',
  210. json)
  211. j.write(subst)
  212. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  213. except Exception as e:
  214. LOGGER.info(f'\n{prefix} export failure: {e}')
  215. def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
  216. try:
  217. check_requirements(('tensorrt',))
  218. import tensorrt as trt
  219. opset = (12, 13)[trt.__version__[0] == '8'] # test on TensorRT 7.x and 8.x
  220. export_onnx(model, im, file, opset, train, False, simplify)
  221. onnx = file.with_suffix('.onnx')
  222. assert onnx.exists(), f'failed to export ONNX file: {onnx}'
  223. LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
  224. f = file.with_suffix('.engine') # TensorRT engine file
  225. logger = trt.Logger(trt.Logger.INFO)
  226. if verbose:
  227. logger.min_severity = trt.Logger.Severity.VERBOSE
  228. builder = trt.Builder(logger)
  229. config = builder.create_builder_config()
  230. config.max_workspace_size = workspace * 1 << 30
  231. flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
  232. network = builder.create_network(flag)
  233. parser = trt.OnnxParser(network, logger)
  234. if not parser.parse_from_file(str(onnx)):
  235. raise RuntimeError(f'failed to load ONNX file: {onnx}')
  236. inputs = [network.get_input(i) for i in range(network.num_inputs)]
  237. outputs = [network.get_output(i) for i in range(network.num_outputs)]
  238. LOGGER.info(f'{prefix} Network Description:')
  239. for inp in inputs:
  240. LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
  241. for out in outputs:
  242. LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
  243. half &= builder.platform_has_fast_fp16
  244. LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}')
  245. if half:
  246. config.set_flag(trt.BuilderFlag.FP16)
  247. with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
  248. t.write(engine.serialize())
  249. LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  250. except Exception as e:
  251. LOGGER.info(f'\n{prefix} export failure: {e}')
  252. @torch.no_grad()
  253. def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
  254. weights=ROOT / 'yolov5s.pt', # weights path
  255. imgsz=(640, 640), # image (height, width)
  256. batch_size=1, # batch size
  257. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  258. include=('torchscript', 'onnx', 'coreml'), # include formats
  259. half=False, # FP16 half-precision export
  260. inplace=False, # set YOLOv5 Detect() inplace=True
  261. train=False, # model.train() mode
  262. optimize=False, # TorchScript: optimize for mobile
  263. int8=False, # CoreML/TF INT8 quantization
  264. dynamic=False, # ONNX/TF: dynamic axes
  265. simplify=False, # ONNX: simplify model
  266. opset=14, # ONNX: opset version
  267. verbose=False, # TensorRT: verbose log
  268. workspace=4, # TensorRT: workspace size (GB)
  269. nms=False, # TF: add NMS to model
  270. agnostic_nms=False, # TF: add agnostic NMS to model
  271. topk_per_class=100, # TF.js NMS: topk per class to keep
  272. topk_all=100, # TF.js NMS: topk for all classes to keep
  273. iou_thres=0.45, # TF.js NMS: IoU threshold
  274. conf_thres=0.25 # TF.js NMS: confidence threshold
  275. ):
  276. t = time.time()
  277. include = [x.lower() for x in include]
  278. tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
  279. imgsz *= 2 if len(imgsz) == 1 else 1 # expand
  280. file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
  281. # Load PyTorch model
  282. device = select_device(device)
  283. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  284. model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
  285. nc, names = model.nc, model.names # number of classes, class names
  286. # Input
  287. gs = int(max(model.stride)) # grid size (max stride)
  288. imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
  289. im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
  290. # Update model
  291. if half:
  292. im, model = im.half(), model.half() # to FP16
  293. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  294. for k, m in model.named_modules():
  295. if isinstance(m, Conv): # assign export-friendly activations
  296. if isinstance(m.act, nn.SiLU):
  297. m.act = SiLU()
  298. elif isinstance(m, Detect):
  299. m.inplace = inplace
  300. m.onnx_dynamic = dynamic
  301. # m.forward = m.forward_export # assign forward (optional)
  302. for _ in range(2):
  303. y = model(im) # dry runs
  304. LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
  305. # Exports
  306. if 'torchscript' in include:
  307. export_torchscript(model, im, file, optimize)
  308. if 'onnx' in include:
  309. export_onnx(model, im, file, opset, train, dynamic, simplify)
  310. if 'engine' in include:
  311. export_engine(model, im, file, train, half, simplify, workspace, verbose)
  312. if 'coreml' in include:
  313. export_coreml(model, im, file)
  314. # TensorFlow Exports
  315. if any(tf_exports):
  316. pb, tflite, tfjs = tf_exports[1:]
  317. assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
  318. model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
  319. agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all,
  320. conf_thres=conf_thres, iou_thres=iou_thres) # keras model
  321. if pb or tfjs: # pb prerequisite to tfjs
  322. export_pb(model, im, file)
  323. if tflite:
  324. export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
  325. if tfjs:
  326. export_tfjs(model, im, file)
  327. # Finish
  328. LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
  329. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  330. f'\nVisualize with https://netron.app')
  331. def parse_opt():
  332. parser = argparse.ArgumentParser()
  333. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  334. parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
  335. parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  336. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  337. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  338. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  339. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  340. parser.add_argument('--train', action='store_true', help='model.train() mode')
  341. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  342. parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  343. parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
  344. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  345. parser.add_argument('--opset', type=int, default=14, help='ONNX: opset version')
  346. parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
  347. parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
  348. parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
  349. parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
  350. parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
  351. parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
  352. parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
  353. parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
  354. parser.add_argument('--include', nargs='+',
  355. default=['torchscript', 'onnx'],
  356. help='available formats are (torchscript, onnx, engine, coreml, saved_model, pb, tflite, tfjs)')
  357. opt = parser.parse_args()
  358. print_args(FILE.stem, opt)
  359. return opt
  360. def main(opt):
  361. for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
  362. run(**vars(opt))
  363. if __name__ == "__main__":
  364. opt = parse_opt()
  365. main(opt)