TensorRT转化代码
Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

190 lines
7.8KB

  1. """Export a YOLOv5 *.pt model to TorchScript, ONNX, CoreML formats
  2. Usage:
  3. $ python path/to/export.py --weights yolov5s.pt --img 640 --batch 1
  4. """
  5. import argparse
  6. import sys
  7. import time
  8. from pathlib import Path
  9. import torch
  10. import torch.nn as nn
  11. from torch.utils.mobile_optimizer import optimize_for_mobile
  12. FILE = Path(__file__).absolute()
  13. sys.path.append(FILE.parents[0].as_posix()) # add yolov5/ to path
  14. from models.common import Conv
  15. from models.yolo import Detect
  16. from models.experimental import attempt_load
  17. from utils.activations import Hardswish, SiLU
  18. from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
  19. from utils.torch_utils import select_device
  20. def export_torchscript(model, img, file, optimize):
  21. # TorchScript model export
  22. prefix = colorstr('TorchScript:')
  23. try:
  24. print(f'\n{prefix} starting export with torch {torch.__version__}...')
  25. f = file.with_suffix('.torchscript.pt')
  26. ts = torch.jit.trace(model, img, strict=False)
  27. (optimize_for_mobile(ts) if optimize else ts).save(f)
  28. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  29. return ts
  30. except Exception as e:
  31. print(f'{prefix} export failure: {e}')
  32. def export_onnx(model, img, file, opset, train, dynamic, simplify):
  33. # ONNX model export
  34. prefix = colorstr('ONNX:')
  35. try:
  36. check_requirements(('onnx', 'onnx-simplifier'))
  37. import onnx
  38. print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  39. f = file.with_suffix('.onnx')
  40. torch.onnx.export(model, img, f, verbose=False, opset_version=opset,
  41. training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  42. do_constant_folding=not train,
  43. input_names=['images'],
  44. output_names=['output'],
  45. dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
  46. 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
  47. } if dynamic else None)
  48. # Checks
  49. model_onnx = onnx.load(f) # load onnx model
  50. onnx.checker.check_model(model_onnx) # check onnx model
  51. # print(onnx.helper.printable_graph(model_onnx.graph)) # print
  52. # Simplify
  53. if simplify:
  54. try:
  55. import onnxsim
  56. print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  57. model_onnx, check = onnxsim.simplify(
  58. model_onnx,
  59. dynamic_input_shape=dynamic,
  60. input_shapes={'images': list(img.shape)} if dynamic else None)
  61. assert check, 'assert check failed'
  62. onnx.save(model_onnx, f)
  63. except Exception as e:
  64. print(f'{prefix} simplifier failure: {e}')
  65. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  66. print(f"{prefix} run --dynamic ONNX model inference with detect.py: 'python detect.py --weights {f}'")
  67. except Exception as e:
  68. print(f'{prefix} export failure: {e}')
  69. def export_coreml(model, img, file):
  70. # CoreML model export
  71. prefix = colorstr('CoreML:')
  72. try:
  73. import coremltools as ct
  74. print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  75. f = file.with_suffix('.mlmodel')
  76. model.train() # CoreML exports should be placed in model.train() mode
  77. ts = torch.jit.trace(model, img, strict=False) # TorchScript model
  78. model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
  79. model.save(f)
  80. print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
  81. except Exception as e:
  82. print(f'\n{prefix} export failure: {e}')
  83. def run(weights='./yolov5s.pt', # weights path
  84. img_size=(640, 640), # image (height, width)
  85. batch_size=1, # batch size
  86. device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
  87. include=('torchscript', 'onnx', 'coreml'), # include formats
  88. half=False, # FP16 half-precision export
  89. inplace=False, # set YOLOv5 Detect() inplace=True
  90. train=False, # model.train() mode
  91. optimize=False, # TorchScript: optimize for mobile
  92. dynamic=False, # ONNX: dynamic axes
  93. simplify=False, # ONNX: simplify model
  94. opset=12, # ONNX: opset version
  95. ):
  96. t = time.time()
  97. include = [x.lower() for x in include]
  98. img_size *= 2 if len(img_size) == 1 else 1 # expand
  99. file = Path(weights)
  100. # Load PyTorch model
  101. device = select_device(device)
  102. assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
  103. model = attempt_load(weights, map_location=device) # load FP32 model
  104. names = model.names
  105. # Input
  106. gs = int(max(model.stride)) # grid size (max stride)
  107. img_size = [check_img_size(x, gs) for x in img_size] # verify img_size are gs-multiples
  108. img = torch.zeros(batch_size, 3, *img_size).to(device) # image size(1,3,320,192) iDetection
  109. # Update model
  110. if half:
  111. img, model = img.half(), model.half() # to FP16
  112. model.train() if train else model.eval() # training mode = no Detect() layer grid construction
  113. for k, m in model.named_modules():
  114. if isinstance(m, Conv): # assign export-friendly activations
  115. if isinstance(m.act, nn.Hardswish):
  116. m.act = Hardswish()
  117. elif isinstance(m.act, nn.SiLU):
  118. m.act = SiLU()
  119. elif isinstance(m, Detect):
  120. m.inplace = inplace
  121. m.onnx_dynamic = dynamic
  122. # m.forward = m.forward_export # assign forward (optional)
  123. for _ in range(2):
  124. y = model(img) # dry runs
  125. print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)")
  126. # Exports
  127. if 'torchscript' in include:
  128. export_torchscript(model, img, file, optimize)
  129. if 'onnx' in include:
  130. export_onnx(model, img, file, opset, train, dynamic, simplify)
  131. if 'coreml' in include:
  132. export_coreml(model, img, file)
  133. # Finish
  134. print(f'\nExport complete ({time.time() - t:.2f}s)'
  135. f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  136. f'\nVisualize with https://netron.app')
  137. def parse_opt():
  138. parser = argparse.ArgumentParser()
  139. parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
  140. parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (height, width)')
  141. parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  142. parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  143. parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
  144. parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  145. parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  146. parser.add_argument('--train', action='store_true', help='model.train() mode')
  147. parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  148. parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
  149. parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
  150. parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
  151. opt = parser.parse_args()
  152. return opt
  153. def main(opt):
  154. set_logging()
  155. print(colorstr('export: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
  156. run(**vars(opt))
  157. if __name__ == "__main__":
  158. opt = parse_opt()
  159. main(opt)