YOLOv5 v6.1发布！支持TensorRT+OpenVINO+TFLite等平台一键导出和部署

https://github.com/ultralytics/yolov5

yolov5 release 6.1版本增加了TensorRT、Edge TPU和OpenVINO的支持，并提供了新的默认单周期线性LR调度器，以128批处理大小的再训练模型。YOLOv5现在正式支持11种不同的权重，不仅可以直接导出，还可以用于推理(detect.py和PyTorch Hub)，以及在导出后对mAP配置文件和速度结果进行验证。

比如，onnx文件的导出：

onnx导出

1重大更新

• TensorRT支持:TensorFlow, Keras, TFLite, TF.js模型导出现在完全集成使用python export.py -include saved_model pb TFLite tfjs

• TensorFlow Edge TPU:新的更小的YOLOv5n(1.9M params)模型低于YOLOv5s(7.5M params)，导出到2.1 MB INT8大小，理想的超轻边缘解决方案。

• OpenVINO支持:YOLOv5 ONNX模型现在兼容OpenCV DNN和ONNX运行。

• Export Benchmarks:使用python utils/ Benchmark.py导出所有YOLOv5格式(mAP和速度)。目前在CPU上运行，未来的更新将实现GPU支持。

• 架构:无更改。

• 超参数:小更改。Yaml LRF从0.2降至0.1。

• 训练:默认学习速率(LR)调度器更新了一个周期的余弦替换为一个周期的线性，以改善结果。

新版模型导出

1、onnx

1.  
   def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr(‘ONNX:’)):
2.  
       # YOLOv5 ONNX export
3.  
       try:
4.  
           check_requirements((‘onnx’,))
5.  
           import onnx
6.  
    
7.  
           LOGGER.info(f‘\n{prefix} starting export with onnx {onnx.__version__}…’)
8.  
           f = file.with_suffix(‘.onnx’)
9.  
    
10.  
            torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
11.  
                              training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
12.  
                              do_constant_folding=not train,
13.  
                              input_names=[‘images’],
14.  
                              output_names=[‘output’],
15.  
                              dynamic_axes={‘images’: {0: ‘batch’, 2: ‘height’, 3: ‘width’},  # shape(1,3,640,640)
16.  
                                            ‘output’: {0: ‘batch’, 1: ‘anchors’}  # shape(1,25200,85)
17.  
                                            } if dynamic else None)
18.  
     
19.  
            # Checks
20.  
            model_onnx = onnx.load(f)  # load onnx model
21.  
            onnx.checker.check_model(model_onnx)  # check onnx model
22.  
            # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print
23.  
     
24.  
            # Simplify
25.  
            if simplify:
26.  
                try:
27.  
                    check_requirements((‘onnx-simplifier’,))
28.  
                    import onnxsim
29.  
     
30.  
                    LOGGER.info(f‘{prefix} simplifying with onnx-simplifier {onnxsim.__version__}…’)
31.  
                    model_onnx, check = onnxsim.simplify(
32.  
                        model_onnx,
33.  
                        dynamic_input_shape=dynamic,
34.  
                        input_shapes={‘images’: list(im.shape)} if dynamic else None)
35.  
                    assert check, ‘assert check failed’
36.  
                    onnx.save(model_onnx, f)
37.  
                except Exception as e:
38.  
                    LOGGER.info(f‘{prefix} simplifier failure: {e}’)
39.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
40.  
            return f
41.  
        except Exception as e:
42.  
            LOGGER.info(f‘{prefix} export failure: {e}’)

2、 openvino

1.  
   def export_openvino(model, im, file, prefix=colorstr(‘OpenVINO:’)):
2.  
       # YOLOv5 OpenVINO export
3.  
       try:
4.  
           check_requirements((‘openvino-dev’,))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
5.  
           import openvino.inference_engine as ie
6.  
    
7.  
           LOGGER.info(f‘\n{prefix} starting export with openvino {ie.__version__}…’)
8.  
           f = str(file).replace(‘.pt’, ‘_openvino_model’ + os.sep)
9.  
    
10.  
            cmd = f“mo –input_model {file.with_suffix(‘.onnx’)} –output_dir {f}”
11.  
            subprocess.check_output(cmd, shell=True)
12.  
     
13.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
14.  
            return f
15.  
        except Exception as e:
16.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

3、 coreml

1.  
   def export_coreml(model, im, file, prefix=colorstr(‘CoreML:’)):
2.  
       # YOLOv5 CoreML export
3.  
       try:
4.  
           check_requirements((‘coremltools’,))
5.  
           import coremltools as ct
6.  
    
7.  
           LOGGER.info(f‘\n{prefix} starting export with coremltools {ct.__version__}…’)
8.  
           f = file.with_suffix(‘.mlmodel’)
9.  
    
10.  
            ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
11.  
            ct_model = ct.convert(ts, inputs=[ct.ImageType(‘image’, shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
12.  
            ct_model.save(f)
13.  
     
14.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
15.  
            return ct_model, f
16.  
        except Exception as e:
17.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)
18.  
            return None, None

4、TensorRT

1.  
   def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr(‘TensorRT:’)):
2.  
       # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
3.  
       try:
4.  
           check_requirements((‘tensorrt’,))
5.  
           import tensorrt as trt
6.  
    
7.  
           if trt.__version__[0] == ‘7’:  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
8.  
               grid = model.model[-1].anchor_grid
9.  
               model.model[-1].anchor_grid = [a[…, :1, :1, :] for a in grid]
10.  
                export_onnx(model, im, file, 12, train, False, simplify)  # opset 12
11.  
                model.model[-1].anchor_grid = grid
12.  
            else:  # TensorRT >= 8
13.  
                check_version(trt.__version__, ‘8.0.0’, hard=True)  # require tensorrt>=8.0.0
14.  
                export_onnx(model, im, file, 13, train, False, simplify)  # opset 13
15.  
            onnx = file.with_suffix(‘.onnx’)
16.  
     
17.  
            LOGGER.info(f‘\n{prefix} starting export with TensorRT {trt.__version__}…’)
18.  
            assert im.device.type != ‘cpu’, ‘export running on CPU but must be on GPU, i.e. `python export.py –device 0`’
19.  
            assert onnx.exists(), f‘failed to export ONNX file: {onnx}’
20.  
            f = file.with_suffix(‘.engine’)  # TensorRT engine file
21.  
            logger = trt.Logger(trt.Logger.INFO)
22.  
            if verbose:
23.  
                logger.min_severity = trt.Logger.Severity.VERBOSE
24.  
     
25.  
            builder = trt.Builder(logger)
26.  
            config = builder.create_builder_config()
27.  
            config.max_workspace_size = workspace * 1 << 30
28.  
     
29.  
            flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
30.  
            network = builder.create_network(flag)
31.  
            parser = trt.OnnxParser(network, logger)
32.  
            if not parser.parse_from_file(str(onnx)):
33.  
                raise RuntimeError(f‘failed to load ONNX file: {onnx}’)
34.  
     
35.  
            inputs = [network.get_input(i) for i in range(network.num_inputs)]
36.  
            outputs = [network.get_output(i) for i in range(network.num_outputs)]
37.  
            LOGGER.info(f‘{prefix} Network Description:’)
38.  
            for inp in inputs:
39.  
                LOGGER.info(f‘{prefix}\tinput “{inp.name}” with shape {inp.shape} and dtype {inp.dtype}’)
40.  
            for out in outputs:
41.  
                LOGGER.info(f‘{prefix}\toutput “{out.name}” with shape {out.shape} and dtype {out.dtype}’)
42.  
     
43.  
            half &= builder.platform_has_fast_fp16
44.  
            LOGGER.info(f‘{prefix} building FP{16 if half else 32} engine in {f}’)
45.  
            if half:
46.  
                config.set_flag(trt.BuilderFlag.FP16)
47.  
            with builder.build_engine(network, config) as engine, open(f, ‘wb’) as t:
48.  
                t.write(engine.serialize())
49.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
50.  
            return f
51.  
        except Exception as e:
52.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

5、TensorFlow

1.  
   def export_pb(keras_model, im, file, prefix=colorstr(‘TensorFlow GraphDef:’)):
2.  
       # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
3.  
       try:
4.  
           import tensorflow as tf
5.  
           from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
6.  
    
7.  
           LOGGER.info(f‘\n{prefix} starting export with tensorflow {tf.__version__}…’)
8.  
           f = file.with_suffix(‘.pb’)
9.  
    
10.  
            m = tf.function(lambda x: keras_model(x))  # full model
11.  
            m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
12.  
            frozen_func = convert_variables_to_constants_v2(m)
13.  
            frozen_func.graph.as_graph_def()
14.  
            tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
15.  
     
16.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
17.  
            return f
18.  
        except Exception as e:
19.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

6、TensorFlow-Lite

1.  
   def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr(‘TensorFlow Lite:’)):
2.  
       # YOLOv5 TensorFlow Lite export
3.  
       try:
4.  
           import tensorflow as tf
5.  
    
6.  
           LOGGER.info(f‘\n{prefix} starting export with tensorflow {tf.__version__}…’)
7.  
           batch_size, ch, *imgsz = list(im.shape)  # BCHW
8.  
           f = str(file).replace(‘.pt’, ‘-fp16.tflite’)
9.  
    
10.  
            converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
11.  
            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
12.  
            converter.target_spec.supported_types = [tf.float16]
13.  
            converter.optimizations = [tf.lite.Optimize.DEFAULT]
14.  
            if int8:
15.  
                from models.tf import representative_dataset_gen
16.  
                dataset = LoadImages(check_dataset(data)[‘train’], img_size=imgsz, auto=False)  # representative data
17.  
                converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
18.  
                converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
19.  
                converter.target_spec.supported_types = []
20.  
                converter.inference_input_type = tf.uint8  # or tf.int8
21.  
                converter.inference_output_type = tf.uint8  # or tf.int8
22.  
                converter.experimental_new_quantizer = False
23.  
                f = str(file).replace(‘.pt’, ‘-int8.tflite’)
24.  
     
25.  
            tflite_model = converter.convert()
26.  
            open(f, “wb”).write(tflite_model)
27.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
28.  
            return f
29.  
        except Exception as e:
30.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

7、Egde TPU

1.  
   def export_edgetpu(keras_model, im, file, prefix=colorstr(‘Edge TPU:’)):
2.  
       # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
3.  
       try:
4.  
           cmd = ‘edgetpu_compiler –version’
5.  
           help_url = ‘https://coral.ai/docs/edgetpu/compiler/’
6.  
           assert platform.system() == ‘Linux’, f‘export only supported on Linux. See {help_url}’
7.  
           if subprocess.run(cmd + ‘ >/dev/null’, shell=True).returncode != 0:
8.  
               LOGGER.info(f‘\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}’)
9.  
               sudo = subprocess.run(‘sudo –version >/dev/null’, shell=True).returncode == 0  # sudo installed on system
10.  
                for c in [‘curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -‘,
11.  
                          ‘echo “deb https://packages.cloud.google.com/apt coral-edgetpu-stable main” | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list’,
12.  
                          ‘sudo apt-get update’,
13.  
                          ‘sudo apt-get install edgetpu-compiler’]:
14.  
                    subprocess.run(c if sudo else c.replace(‘sudo ‘, ”), shell=True, check=True)
15.  
            ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
16.  
     
17.  
            LOGGER.info(f‘\n{prefix} starting export with Edge TPU compiler {ver}…’)
18.  
            f = str(file).replace(‘.pt’, ‘-int8_edgetpu.tflite’)  # Edge TPU model
19.  
            f_tfl = str(file).replace(‘.pt’, ‘-int8.tflite’)  # TFLite model
20.  
     
21.  
            cmd = f“edgetpu_compiler -s {f_tfl}”
22.  
            subprocess.run(cmd, shell=True, check=True)
23.  
     
24.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
25.  
            return f
26.  
        except Exception as e:
27.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

8、TensorFlow.js

1.  
   def export_tfjs(keras_model, im, file, prefix=colorstr(‘TensorFlow.js:’)):
2.  
       # YOLOv5 TensorFlow.js export
3.  
       try:
4.  
           check_requirements((‘tensorflowjs’,))
5.  
           import re
6.  
    
7.  
           import tensorflowjs as tfjs
8.  
    
9.  
           LOGGER.info(f‘\n{prefix} starting export with tensorflowjs {tfjs.__version__}…’)
10.  
            f = str(file).replace(‘.pt’, ‘_web_model’)  # js dir
11.  
            f_pb = file.with_suffix(‘.pb’)  # *.pb path
12.  
            f_json = f + ‘/model.json’  # *.json path
13.  
     
14.  
            cmd = f‘tensorflowjs_converter –input_format=tf_frozen_model ‘ \
15.  
                  f‘–output_node_names=”Identity,Identity_1,Identity_2,Identity_3″ {f_pb} {f}’
16.  
            subprocess.run(cmd, shell=True)
17.  
     
18.  
            json = open(f_json).read()
19.  
            with open(f_json, ‘w’) as j:  # sort JSON Identity_* in ascending order
20.  
                subst = re.sub(
21.  
                    r‘{“outputs”: {“Identity.?.?”: {“name”: “Identity.?.?”}, ‘
22.  
                    r‘”Identity.?.?”: {“name”: “Identity.?.?”}, ‘
23.  
                    r‘”Identity.?.?”: {“name”: “Identity.?.?”}, ‘
24.  
                    r‘”Identity.?.?”: {“name”: “Identity.?.?”}}}’,
25.  
                    r‘{“outputs”: {“Identity”: {“name”: “Identity”}, ‘
26.  
                    r‘”Identity_1″: {“name”: “Identity_1”}, ‘
27.  
                    r‘”Identity_2″: {“name”: “Identity_2”}, ‘
28.  
                    r‘”Identity_3″: {“name”: “Identity_3”}}}’,
29.  
                    json)
30.  
                j.write(subst)
31.  
     
32.  
            LOGGER.info(f‘{prefix} export success, saved as {f} ({file_size(f):.1f} MB)’)
33.  
            return f
34.  
        except Exception as e:
35.  
            LOGGER.info(f‘\n{prefix} export failure: {e}’)

新版检测推理

1.  
   python path/to/detect.py –weights yolov5s.pt                 # PyTorch
2.  
                                      yolov5s.torchscript        # TorchScript
3.  
                                      yolov5s.onnx               # ONNX Runtime or OpenCV DNN with –dnn
4.  
                                      yolov5s.xml                # OpenVINO
5.  
                                      yolov5s.engine             # TensorRT
6.  
                                      yolov5s.mlmodel            # CoreML (MacOS-only)
7.  
                                      yolov5s_saved_model        # TensorFlow SavedModel
8.  
                                      yolov5s.pb                 # TensorFlow GraphDef
9.  
                                      yolov5s.tflite             # TensorFlow Lite
10.  
                                       yolov5s_edgetpu.tflite     # TensorFlow Edge TPU

2最新结果

3与6.0版本的精度对比

1.  
   ICCV和CVPR 2021论文和代码下载
2.  
    
3.  
   后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集
4.  
    
5.  
   后台回复：ICCV2021，即可下载ICCV 2021论文和代码开源的论文合集
6.  
    
7.  
   后台回复：Transformer综述，即可下载最新的3篇Transformer综述PDF
8.  
   CVer-目标检测交流群成立
9.  
   扫码添加CVer助手，可申请加入CVer-目标检测 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
10.  
    一定要备注：研究方向+地点+学校/公司+昵称（如目标检测+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群
11.  
     
12.  
    ▲长按加小助手微信，进交流群
13.  
    CVer学术交流群（知识星球）来了！想要了解最新最快最好的CV/DL/ML论文速递、优质开源项目、学习教程和实战训练等资料，欢迎扫描下方二维码，加入CVer学术交流群，已汇集数千人！
14.  
     
15.  
    ▲扫码进群
16.  
    ▲点击上方卡片，关注CVer公众号
17.  
     
18.  
    整理不易，请点赞和在看

 转自https://blog.csdn.net/amusi1994/article/details/123102969