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.

338 lignes
15KB

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