YOLOV5 v6.1更新 | TensorRT+TPU+OpenVINO+TFJS+TFLite等平台一键导出和部署

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

        
          
def export\_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):  
    # YOLOv5 ONNX export  
    try:  
        check_requirements(('onnx',))  
        import onnx  
  
        LOGGER.info(f'\n{prefix} starting export with onnx {onnx.\_\_version\_\_}...')  
        f = file.with_suffix('.onnx')  
  
        torch.onnx.export(model, im, f, verbose=False, opset_version=opset,  
                          training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,  
                          do_constant_folding=not train,  
                          input_names=['images'],  
                          output_names=['output'],  
                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)  
                                        'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)  
                                        } if dynamic else None)  
  
        # Checks  
        model_onnx = onnx.load(f)  # load onnx model  
        onnx.checker.check_model(model_onnx)  # check onnx model  
        # LOGGER.info(onnx.helper.printable\_graph(model\_onnx.graph))  # print  
  
        # Simplify  
        if simplify:  
            try:  
                check_requirements(('onnx-simplifier',))  
                import onnxsim  
  
                LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.\_\_version\_\_}...')  
                model_onnx, check = onnxsim.simplify(  
                    model_onnx,  
                    dynamic_input_shape=dynamic,  
                    input_shapes={'images': list(im.shape)} if dynamic else None)  
                assert check, 'assert check failed'  
                onnx.save(model_onnx, f)  
            except Exception as e:  
                LOGGER.info(f'{prefix} simplifier failure: {e}')  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'{prefix} export failure: {e}')  

      

2、 openvino

        
          
def export\_openvino(model, im, file, prefix=colorstr('OpenVINO:')):  
    # YOLOv5 OpenVINO export  
    try:  
        check_requirements(('openvino-dev',))  # requires openvino-dev: https://pypi.org/project/openvino-dev/  
        import openvino.inference_engine as ie  
  
        LOGGER.info(f'\n{prefix} starting export with openvino {ie.\_\_version\_\_}...')  
        f = str(file).replace('.pt', '\_openvino\_model' + os.sep)  
  
        cmd = f"mo --input\_model {file.with\_suffix('.onnx')} --output\_dir {f}"  
        subprocess.check_output(cmd, shell=True)  
  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

3、 coreml

        
          
def export\_coreml(model, im, file, prefix=colorstr('CoreML:')):  
    # YOLOv5 CoreML export  
    try:  
        check_requirements(('coremltools',))  
        import coremltools as ct  
  
        LOGGER.info(f'\n{prefix} starting export with coremltools {ct.\_\_version\_\_}...')  
        f = file.with_suffix('.mlmodel')  
  
        ts = torch.jit.trace(model, im, strict=False)  # TorchScript model  
        ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])  
        ct_model.save(f)  
  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return ct_model, f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  
        return None, None  
  

      

4、TensorRT

        
          
def export\_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):  
    # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt  
    try:  
        check_requirements(('tensorrt',))  
        import tensorrt as trt  
  
        if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012  
            grid = model.model[-1].anchor_grid  
            model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]  
            export_onnx(model, im, file, 12, train, False, simplify)  # opset 12  
            model.model[-1].anchor_grid = grid  
        else:  # TensorRT >= 8  
            check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0  
            export_onnx(model, im, file, 13, train, False, simplify)  # opset 13  
        onnx = file.with_suffix('.onnx')  
  
        LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.\_\_version\_\_}...')  
        assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'  
        assert onnx.exists(), f'failed to export ONNX file: {onnx}'  
        f = file.with_suffix('.engine')  # TensorRT engine file  
        logger = trt.Logger(trt.Logger.INFO)  
        if verbose:  
            logger.min_severity = trt.Logger.Severity.VERBOSE  
  
        builder = trt.Builder(logger)  
        config = builder.create_builder_config()  
        config.max_workspace_size = workspace * 1 << 30  
  
        flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))  
        network = builder.create_network(flag)  
        parser = trt.OnnxParser(network, logger)  
        if not parser.parse_from_file(str(onnx)):  
            raise RuntimeError(f'failed to load ONNX file: {onnx}')  
  
        inputs = [network.get_input(i) for i in range(network.num_inputs)]  
        outputs = [network.get_output(i) for i in range(network.num_outputs)]  
        LOGGER.info(f'{prefix} Network Description:')  
        for inp in inputs:  
            LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')  
        for out in outputs:  
            LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')  
  
        half &= builder.platform_has_fast_fp16  
        LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}')  
        if half:  
            config.set_flag(trt.BuilderFlag.FP16)  
        with builder.build_engine(network, config) as engine, open(f, 'wb') as t:  
            t.write(engine.serialize())  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

5、TensorFlow

        
          
def export\_pb(keras\_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):  
    # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen\_Graph\_TensorFlow  
    try:  
        import tensorflow as tf  
        from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2  
  
        LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.\_\_version\_\_}...')  
        f = file.with_suffix('.pb')  
  
        m = tf.function(lambda x: keras_model(x))  # full model  
        m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))  
        frozen_func = convert_variables_to_constants_v2(m)  
        frozen_func.graph.as_graph_def()  
        tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)  
  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

6、TensorFlow-Lite

        
          
def export\_tflite(keras\_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):  
    # YOLOv5 TensorFlow Lite export  
    try:  
        import tensorflow as tf  
  
        LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.\_\_version\_\_}...')  
        batch_size, ch, *imgsz = list(im.shape)  # BCHW  
        f = str(file).replace('.pt', '-fp16.tflite')  
  
        converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)  
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]  
        converter.target_spec.supported_types = [tf.float16]  
        converter.optimizations = [tf.lite.Optimize.DEFAULT]  
        if int8:  
            from models.tf import representative_dataset_gen  
            dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False)  # representative data  
            converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)  
            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]  
            converter.target_spec.supported_types = []  
            converter.inference_input_type = tf.uint8  # or tf.int8  
            converter.inference_output_type = tf.uint8  # or tf.int8  
            converter.experimental_new_quantizer = False  
            f = str(file).replace('.pt', '-int8.tflite')  
  
        tflite_model = converter.convert()  
        open(f, "wb").write(tflite_model)  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

7、Egde TPU

        
          
def export\_edgetpu(keras\_model, im, file, prefix=colorstr('Edge TPU:')):  
    # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/  
    try:  
        cmd = 'edgetpu\_compiler --version'  
        help_url = 'https://coral.ai/docs/edgetpu/compiler/'  
        assert platform.system() == 'Linux', f'export only supported on Linux. See {help\_url}'  
        if subprocess.run(cmd + ' >/dev/null', shell=True).returncode != 0:  
            LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help\_url}')  
            sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0  # sudo installed on system  
            for c in ['curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',  
                      'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',  
                      'sudo apt-get update',  
                      'sudo apt-get install edgetpu-compiler']:  
                subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)  
        ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]  
  
        LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')  
        f = str(file).replace('.pt', '-int8\_edgetpu.tflite')  # Edge TPU model  
        f_tfl = str(file).replace('.pt', '-int8.tflite')  # TFLite model  
  
        cmd = f"edgetpu\_compiler -s {f\_tfl}"  
        subprocess.run(cmd, shell=True, check=True)  
  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

8、TensorFlow.js

        
          
def export\_tfjs(keras\_model, im, file, prefix=colorstr('TensorFlow.js:')):  
    # YOLOv5 TensorFlow.js export  
    try:  
        check_requirements(('tensorflowjs',))  
        import re  
  
        import tensorflowjs as tfjs  
  
        LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.\_\_version\_\_}...')  
        f = str(file).replace('.pt', '\_web\_model')  # js dir  
        f_pb = file.with_suffix('.pb')  # *.pb path  
        f_json = f + '/model.json'  # *.json path  
  
        cmd = f'tensorflowjs\_converter --input\_format=tf\_frozen\_model ' \  
              f'--output\_node\_names="Identity,Identity\_1,Identity\_2,Identity\_3" {f\_pb} {f}'  
        subprocess.run(cmd, shell=True)  
  
        json = open(f_json).read()  
        with open(f_json, 'w') as j:  # sort JSON Identity\_* in ascending order  
            subst = re.sub(  
                r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '  
                r'"Identity.?.?": {"name": "Identity.?.?"}, '  
                r'"Identity.?.?": {"name": "Identity.?.?"}, '  
                r'"Identity.?.?": {"name": "Identity.?.?"}}}',  
                r'{"outputs": {"Identity": {"name": "Identity"}, '  
                r'"Identity\_1": {"name": "Identity\_1"}, '  
                r'"Identity\_2": {"name": "Identity\_2"}, '  
                r'"Identity\_3": {"name": "Identity\_3"}}}',  
                json)  
            j.write(subst)  
  
        LOGGER.info(f'{prefix} export success, saved as {f} ({file\_size(f):.1f} MB)')  
        return f  
    except Exception as e:  
        LOGGER.info(f'\n{prefix} export failure: {e}')  

      

新版检测推理

        
          
python path/to/detect.py --weights yolov5s.pt                 # PyTorch  
                                   yolov5s.torchscript        # TorchScript  
                                   yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn  
                                   yolov5s.xml                # OpenVINO  
                                   yolov5s.engine             # TensorRT  
                                   yolov5s.mlmodel            # CoreML (MacOS-only)  
                                   yolov5s_saved_model        # TensorFlow SavedModel  
                                   yolov5s.pb                 # TensorFlow GraphDef  
                                   yolov5s.tflite             # TensorFlow Lite  
                                   yolov5s_edgetpu.tflite     # TensorFlow Edge TPU  

      

2

最新结果

3

与6.0版本的精度对比

4

参考

[1].https://github.com/ultralytics/yolov5

5

推荐阅读

往期推荐

超越ConvNeXt | 大道至简，VAN用普通卷积，登顶Backbone性能巅峰（附代码解读）

超级干货 | 用万字文章总结25种正则化方法（值得收藏）

小目标Trick | Detectron2、MMDetection、YOLOv5都通用的小目标检测解决方案

ShiftViT用Swin Transformer的精度跑赢ResNet的速度，论述ViT的成功不在注意力！

阿里提出QuadTree Transformer | 最轻、最强的Vision Transformer Backbone

MoA-Transformer | Swin-Transformer应该如何更好地引入全局信息？

致敬ATSS | Dynamic ATSS再造ATSS辉煌！！！

长按扫描下方二维码添加小助手并加入交流群， 群里博士大佬云集， 每日讨论话题有目标检测、语义分 割、 超分辨率、模型部署、数学基础知识、算法面试题分享的等等内容，当然也少不了搬砖人的扯犊子

长按扫描下方二维码添加小助手。

可以一起讨论遇到的问题

声明：转载请说明出处

扫描下方二维码关注【 集智书童 】公众号，获取更多实践项目源码和论文解读，非常期待你我的相遇，让我们以梦为马，砥砺前行！