YOLOV7算法(五)pth/pt转onnx学习记录

这篇具有很好参考价值的文章主要介绍了YOLOV7算法(五)pth/pt转onnx学习记录。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

输入指令

python export.py --weights /kaxier01/projects/FAS/yolov7/weights/yolov7.pt --grid --end2end --simplify --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 --img-size 640 640 --max-wh 640

export.py代码学习

import argparse
import sys
import time
import warnings

sys.path.append('./')  # to run '$ python *.py' files in subdirectories

import torch
import torch.nn as nn
from torch.utils.mobile_optimizer import optimize_for_mobile

import models
from models.experimental import attempt_load, End2End
from utils.activations import Hardswish, SiLU
from utils.general import set_logging, check_img_size
from utils.torch_utils import select_device
from utils.add_nms import RegisterNMS
import sys
import warnings
warnings.filterwarnings('ignore')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='/kaxier01/projects/FAS/yolov7/weights/yolov7.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
    parser.add_argument('--dynamic-batch', action='store_true', help='dynamic batch onnx for tensorrt and onnx-runtime')
    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
    parser.add_argument('--end2end', action='store_true', help='export end2end onnx')
    parser.add_argument('--max-wh', type=int, default=None, help='None for tensorrt nms, int value for onnx-runtime nms')
    parser.add_argument('--topk-all', type=int, default=100, help='topk objects for every images')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='iou threshold for NMS')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='conf threshold for NMS')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--simplify', action='store_true', help='simplify onnx model')
    parser.add_argument('--include-nms', action='store_true', help='export end2end onnx')
    parser.add_argument('--fp16', action='store_true', help='CoreML FP16 half-precision export')
    parser.add_argument('--int8', action='store_true', help='CoreML INT8 quantization')
    opt = parser.parse_args()
    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # opt.img_size=[640, 640]
    opt.dynamic = opt.dynamic and not opt.end2end  # False
    opt.dynamic = False if opt.dynamic_batch else opt.dynamic  # False
    print(opt)
    set_logging()
    t = time.time()

    # Load PyTorch model
    device = select_device(opt.device)  # device='cpu'
    model = attempt_load(opt.weights, map_location=device)  # load FP32 model
    labels = model.names
    '''
    labels=
    ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 
     'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 
     'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 
     'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 
     'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 
     'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 
     'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 
     'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 
     'scissors', 'teddy bear', 'hair drier', 'toothbrush']
    '''

    # Checks
    gs = int(max(model.stride))  # grid size (max stride), gs=32
    opt.img_size = [check_img_size(x, gs) for x in opt.img_size]  # verify img_size are gs-multiples, opt.img_size=[640, 640]

    # Input
    img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)  # image size(1,3,320,192) iDetection, img.shape=torch.Size([1, 3, 640, 640])

    # Update model
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
    
    model.model[-1].export = not opt.grid  # set Detect() layer grid export, model.model[-1].export=False
    y = model(img)  # dry run
    if opt.include_nms:
        model.model[-1].include_nms = True
        y = None

    # TorchScript export
    try:
        print('\nStarting TorchScript export with torch %s...' % torch.__version__)
        f = opt.weights.replace('.pt', '.torchscript.pt')  # f='/kaxier01/projects/FAS/yolov7/weights/yolov7.torchscript.pt'
        ts = torch.jit.trace(model, img, strict=False)
        ts.save(f)  # .torchscript.pt模型可以不依赖于python而直接在c++等环境中运行
        print('TorchScript export success, saved as %s' % f)
    except Exception as e:
        print('TorchScript export failure: %s' % e)

    # CoreML export
    try:
        import coremltools as ct

        print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
        # convert model from torchscript and apply pixel scaling as per detect.py
        ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
        bits, mode = (8, 'kmeans_lut') if opt.int8 else (16, 'linear') if opt.fp16 else (32, None)
        if bits < 32:
            if sys.platform.lower() == 'darwin':  # quantization only supported on macOS
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
                    ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
            else:
                print('quantization only supported on macOS, skipping...')

        f = opt.weights.replace('.pt', '.mlmodel')  # f='/kaxier01/projects/FAS/yolov7/weights/yolov7.mlmodel'
        ct_model.save(f)  # .mlmodel可部署到IOS端
        print('CoreML export success, saved as %s' % f)
    except Exception as e:
        print('CoreML export failure: %s' % e)
                     
    # TorchScript-Lite export
    try:
        print('\nStarting TorchScript-Lite export with torch %s...' % torch.__version__)
        f = opt.weights.replace('.pt', '.torchscript.ptl')  # f='/kaxier01/projects/FAS/yolov7/weights/yolov7.torchscript.ptl'
        tsl = torch.jit.trace(model, img, strict=False)
        tsl = optimize_for_mobile(tsl)
        tsl._save_for_lite_interpreter(f)  # .torchscript.ptl模型可部署到Android端
        print('TorchScript-Lite export success, saved as %s' % f)
    except Exception as e:
        print('TorchScript-Lite export failure: %s' % e)

    # ONNX export
    try:
        import onnx

        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = opt.weights.replace('.pt', '.onnx')  # f='/kaxier01/projects/FAS/yolov7/weights/yolov7.onnx'
        model.eval()
        output_names = ['classes', 'boxes'] if y is None else ['output']  # output_names=['output']
        dynamic_axes = None
        if opt.dynamic:
            dynamic_axes = {'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
             'output': {0: 'batch', 2: 'y', 3: 'x'}}
        if opt.dynamic_batch:
            opt.batch_size = 'batch'
            dynamic_axes = {
                'images': {
                    0: 'batch',
                }, }
            if opt.end2end and opt.max_wh is None:
                output_axes = {
                    'num_dets': {0: 'batch'},
                    'det_boxes': {0: 'batch'},
                    'det_scores': {0: 'batch'},
                    'det_classes': {0: 'batch'},
                }
            else:
                output_axes = {
                    'output': {0: 'batch'},
                }
            dynamic_axes.update(output_axes)
        if opt.grid:
            if opt.end2end:
                print('\nStarting export end2end onnx model for %s...' % 'TensorRT' if opt.max_wh is None else 'onnxruntime')
                model = End2End(model,opt.topk_all,opt.iou_thres,opt.conf_thres,opt.max_wh,device,len(labels))
                if opt.end2end and opt.max_wh is None:
                    output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']
                    shapes = [opt.batch_size, 1, opt.batch_size, opt.topk_all, 4,
                              opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
                else:
                    output_names = ['output']
            else:
                model.model[-1].concat = True

        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                          output_names=output_names,
                          dynamic_axes=dynamic_axes)

        # Checks
        onnx_model = onnx.load(f)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model

        if opt.end2end and opt.max_wh is None:
            for i in onnx_model.graph.output:
                for j in i.type.tensor_type.shape.dim:
                    j.dim_param = str(shapes.pop(0))

        # print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model

        if opt.simplify:
            try:
                import onnxsim

                print('\nStarting to simplify ONNX...')
                onnx_model, check = onnxsim.simplify(onnx_model)  # 简化模型
                assert check, 'assert check failed'
            except Exception as e:
                print(f'Simplifier failure: {e}')

        # print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
        onnx.save(onnx_model,f)
        print('ONNX export success, saved as %s' % f)

        if opt.include_nms:
            print('Registering NMS plugin for ONNX...')
            mo = RegisterNMS(f)
            mo.register_nms()
            mo.save(f)

    except Exception as e:
        print('ONNX export failure: %s' % e)

    # Finish
    print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))

如果遇到

CoreML export failure: Core ML only supports tensors with rank <= 5. Layer "model.105.anchor_grid", with type "const", outputs a rank 6 tensor.

则把输入指令改成

python export.py --weights /kaxier01/projects/FAS/yolov7/weights/yolov7.pt --end2end --simplify --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 --img-size 640 640 --max-wh 640

yolov7.onnx网络结构图文章来源地址https://www.toymoban.com/news/detail-508695.html

到了这里,关于YOLOV7算法(五)pth/pt转onnx学习记录的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • opencv调用yolov7 yolov7 c++ yolov7转onnx opencv调用yolov7 onnx 一、YOLOV7主要贡献:

            主要是现有的一些trick的集合以及模块重参化和动态标签分配策略,最终在 5 FPS 到 160 FPS 范围内的速度和准确度都超过了所有已知的目标检测器。         当前目标检测主要的优化方向:更快更强的网络架构;更有效的特征集成方法;更准确的检测方法;更精

    2023年04月22日
    浏览(41)
  • YOLOV5(二):将pt转为onnx模型并用opencv部署

    yolov5s 6.0自带export.py程序可将.pt转为.onnx等,只需配置需要的环境即可。 1. 安装环境 报错:NVIDIA-tensorrt安装失败! 解决:从源码安装TensorRt: ①安装CUDNN和CudaToolKit等GPU配置 ②从官网下载需要的rt版本:https://developer.nvidia.com/nvidia-tensorrt-8x-download ③解压后,将lib文件夹添加到

    2024年02月10日
    浏览(44)
  • 如何在windows系统下将yolov5的pt模型导出为onnx模型

    最近在做本科毕业设计,要求是在树莓派上部署yolo算法来实现火灾检测,在网上查了很多资料,最后选择用yolov5s模型先试着在树莓派上部署,看下效果如何,由于从大佬那里拿到了yolov5火灾检测模型,但想要将它移植到树莓派上第一步要把pt模型转换成onnx模型,原因我想大

    2023年04月12日
    浏览(45)
  • YOLOV7学习记录之原理+代码介绍

    博主计划做一个目标检测跟踪项目,考虑使用YOLO系列模型来作为目标检测器,如今YOLO项目已经更新到了YOLOV7版本,因此便来学习一下相关原理,完成相关实验工作。 论文链接:https://arxiv.org/abs/2207.02696 YOLOv7是 YOLO 系列最新推出的YOLO 结构,在 5 帧/秒到 160 帧/秒范围内,其速

    2024年02月03日
    浏览(53)
  • yolov7使用onnx推理(带&不带NMS)

    官方代码: GitHub - WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors         上个月,官方放了个使用onnx推理的ipynb文件,过了几天上去看,官方又给删了,不知道是不是要更新波大的,还好手快保存了一份,这个可以

    2024年02月16日
    浏览(40)
  • yolov5训练自己的pt文件,转onnx,再转成rknn,到RK3588开发板运行测试

    yolov5训练好自己的模型,例如训练完后,名称为best.pt,路径为runs/exp/weights/best.pt。 采用detect.py文件验证best.pt可以正常检测目标,再进行下一步工作。 修改utils/yolo.py文件的后处理部分,将class Detect(nn.Module) 类的子函数forward由 修改为: 注意:训练和检测的时候,yolo.py文件应

    2024年02月01日
    浏览(58)
  • 【深度学习】pytorch pth模型转为onnx模型后出现冗余节点“identity”,onnx模型的冗余节点“identity”

    onnx模型的冗余节点“identity”如下图。 首先,确保您已经安装了onnx-simplifier库: 然后,您可以按照以下方式使用onnx-simplifier库: 通过这个过程,onnx-simplifier库将会检测和移除不必要的\\\"identity\\\"节点,从而减少模型中的冗余。 请注意,使用onnx-simplifier库可能会改变模型的计算

    2024年02月09日
    浏览(45)
  • 将pytorch的pth文件固化为pt文件

    说明 我参考了一个开源的人像语义分割项目mobile_phone_human_matting,这个项目提供了预训练模型,我想要将该模型固化,然后转换格式后在嵌入式端使用。 该项目保存模型的代码如下: 转换代码 上面代码保存了state_dict, 所以保存的文件中是不含模型结构的,固化时需要从代码

    2024年02月10日
    浏览(44)
  • pth转onnx,onnx转tflite,亲测有效

            pth模型转onnx模型采用torch框架现有转换函数:torch.onnx.export(),函数使用过程中注意要设置输入输出name,以及将batch维度设置为动态,以便后续onnx转tflite。         具体转换代码如下:         网上其他转换方案,虽然可以成功转换,但是存在一个问题

    2024年02月13日
    浏览(70)
  • YOLOv7-tracker 目标追踪 输入视频帧

    参考项目:https://github.com/JackWoo0831/Yolov7-tracker/tree/master github链接:https://github.com/Whiffe/Yolov7-tracker 码云链接:https://gitee.com/YFwinston/Yolov7-tracker b站视频:https://www.bilibili.com/video/BV1Rh4y1K7zt/ 平台:AutoDL 选择镜像: PyTorch 1.7.0 Python 3.8(ubuntu18.04) Cuda 11.0 通过github下载 或者通过码云

    2024年02月11日
    浏览(31)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包