yolact导出onnx

github上有yolact-onnx仓库可以导出不带有nms和两个分支的矩阵相乘的部分,但是无法导出带有nms的部分。

一、导出的代码

注意opset版本最低要求14, torch.onnx.export的输入要求是真实图片,否则后续推理会报错。

python 复制代码
import torch
import cv2

from yolact import Yolact

def export_onnx_model(saved_onnx_model):
    """
    将模型导出为onnx格式, opset版本最低要设置14, 11的话有个算子不能导出
    """
    device = torch.device('cpu')
    
    net = Yolact()
    net.load_weights('weights/yolact_base_54_800000.pth')
    net.to(device)
    net.eval()
    
    img = cv2.imread('images/test/4.jpg')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (550, 550))
    img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float().to(device)
    # img = torch.randn(1, 3, 550, 550).to(device)

    torch.onnx.export(net, img, saved_onnx_model, verbose=True, opset_version=17)
    
export_onnx_model('yolact.onnx')

二、Bug及解决

  1. FPN
python 复制代码
RuntimeError: Tried to trace <__torch__.yolact.FPN object at 0x6db3f50> but it is not part of the active trace. Modules that are called during a trace must be registered as submodules of the thing being traced.

解决:在yolact.py 第25行将 use_jit 设置为False。

python 复制代码
use_jit = torch.cuda.device_count() <= 1
use_jit = False   
if not use_jit:
    print('Multiple GPUs detected! Turning off JIT.')
  1. numpy
python 复制代码
RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.

解决:detection.py第208行改为如下:

python 复制代码
# preds = torch.cat([boxes[conf_mask], cls_scores[:, None]], dim=1).cpu().numpy()
preds = torch.cat([boxes[conf_mask], cls_scores[:, None]], dim=1).cpu().detach().numpy()

第30行将use_fast_nms改为True:

python 复制代码
self.use_fast_nms = True
# self.use_fast_nms = False
  1. tupels
python 复制代码
RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: Yolact

解决:将detection.py 第76行改为如下。这里的net是后处理用来eval_mask的,但是那个if语句是False,相当于返回去也没用上,这里直接不返回也没关系。不然输出的tuple无法转为onnx。

python 复制代码
# out.append({'detection': result, 'net': net})
out.append(result)

output_utils.py中注释掉net

python 复制代码
dets = det_output[batch_idx]
# net = dets['net']
dets = dets['detection']
...
# if cfg.use_maskiou:
#     with timer.env('maskiou_net'):                
#         with torch.no_grad():
#             maskiou_p = net.maskiou_net(masks.unsqueeze(1))
#             maskiou_p = torch.gather(maskiou_p, dim=1, index=classes.unsqueeze(1)).squeeze(1)
#             if cfg.rescore_mask:
#                 if cfg.rescore_bbox:
#                     scores = scores * maskiou_p
#                 else:
#                     scores = [scores, scores * maskiou_p]
  1. output_utils.py

第74行增加这一句,输出增加 'priors',不然后续推理出错:

python 复制代码
if result is not None and proto_data is not None:
result['proto'] = proto_data[batch_idx]
result['priors'] = prior_data[batch_idx]   # add, important

第62行注释掉,这里默认为False,不会执行。只是为了后面我的验证代码能够正常运行:

python 复制代码
# Test flag, do not upvote
# if cfg.mask_proto_debug:
#     np.save('scripts/proto.npy', proto_data.cpu().numpy())
        
# if visualize_lincomb:
#     display_lincomb(proto_data, masks)
  1. box_utils.py
    (很重要)这里要将@torch.jit.script注释掉,否则到处的结果是错误的:
python 复制代码
# @torch.jit.script
def decode(loc, priors, use_yolo_regressors:bool=False):

至此可以导出。

三、验证

注意修改图片路径:

python 复制代码
import torch
import onnx
import os
import cv2
import torch
import argparse
from data import COCODetection, get_label_map, MEANS, COLORS
# from eval import parse_args
from eval import args
from layers.output_utils import postprocess
import onnxruntime as rt
import thop
from torch.profiler import profile, record_function, ProfilerActivity
from yolact import Yolact
from torch.utils.data import Dataset
from utils.augmentations import BaseTransform, FastBaseTransform, Resize
from layers import Detect
from collections import defaultdict
from data import cfg
from utils import timer

def prep_display(dets_out, img, h, w, undo_transform=True, class_color=False, mask_alpha=0.45, fps_str=''):
    """
    Note: If undo_transform=False then im_h and im_w are allowed to be None.
    """
    # args =parse_args()
    img_gpu = img / 255.0
    h, w, _ = img.shape

    # 后处理 w, h = 612 612
    with timer.env('Post'):
        t = postprocess(dets_out, w, h, visualize_lincomb = args.display_lincomb,
                                    crop_masks        = args.crop,
                                    score_threshold   = args.score_threshold)

    idx = t[1].argsort(0, descending=True)[:args.top_k]  # 012345   args.top_k=5?
    
    # if cfg.eval_mask_branch:
        # Masks are drawn on the GPU, so don't copy
    # masks = t[3][idx]
    classes, scores, boxes, masks = [x[idx].cpu().numpy() for x in t[:4]]  # 5,4    最终后处理的结果
    masks = torch.tensor(masks)

    num_dets_to_consider = min(args.top_k, classes.shape[0])  # 指定要检测的最大目标数 vs 检测出来的目标个数,取最小值
    for j in range(num_dets_to_consider):
        if scores[j] < args.score_threshold:
            num_dets_to_consider = j
            break

    # Quick and dirty lambda for selecting the color for a particular index
    # Also keeps track of a per-gpu color cache for maximum speed
    def get_color(j, on_gpu=None):
        global color_cache
        color_idx = (classes[j] * 5 if class_color else j * 5) % len(COLORS)
        
        if on_gpu is not None and color_idx in color_cache[on_gpu]:
            return color_cache[on_gpu][color_idx]
        else:
            color = COLORS[color_idx]
            if not undo_transform:
                # The image might come in as RGB or BRG, depending
                color = (color[2], color[1], color[0])
            if on_gpu is not None:
                color = torch.Tensor(color).to(on_gpu).float() / 255.
                color_cache[on_gpu][color_idx] = color
            return color

    # First, draw the masks on the GPU where we can do it really fast
    # Beware: very fast but possibly unintelligible mask-drawing code ahead
    # I wish I had access to OpenGL or Vulkan but alas, I guess Pytorch tensor operations will have to suffice
    if args.display_masks and num_dets_to_consider > 0:
        # After this, mask is of size [num_dets, h, w, 1]
        masks = masks[:num_dets_to_consider, :, :, None]
        
        # Prepare the RGB images for each mask given their color (size [num_dets, h, w, 1])
        colors = torch.cat([torch.Tensor(get_color(j, on_gpu=img_gpu.device.index)).view(1, 1, 1, 3) for j in range(num_dets_to_consider)], dim=0)
        masks_color = masks.repeat(1, 1, 1, 3) * colors * mask_alpha  # 3,1,1,3 -->3,h,w,3

        # This is 1 everywhere except for 1-mask_alpha where the mask is
        inv_alph_masks = masks * (-mask_alpha) + 1
        
        # I did the math for this on pen and paper. This whole block should be equivalent to:
        #    for j in range(num_dets_to_consider):
        #        img_gpu = img_gpu * inv_alph_masks[j] + masks_color[j]
        masks_color_summand = masks_color[0]
        if num_dets_to_consider > 1:
            inv_alph_cumul = inv_alph_masks[:(num_dets_to_consider-1)].cumprod(dim=0)
            masks_color_cumul = masks_color[1:] * inv_alph_cumul
            masks_color_summand += masks_color_cumul.sum(dim=0)

        img_gpu = img_gpu * inv_alph_masks.prod(dim=0) + masks_color_summand
    
    # Then draw the stuff that needs to be done on the cpu
    # Note, make sure this is a uint8 tensor or opencv will not anti alias text for whatever reason
    img_numpy = (img_gpu * 255).byte().cpu().numpy()
    
    if num_dets_to_consider == 0:
        return img_numpy

    if args.display_text or args.display_bboxes:
        for j in reversed(range(num_dets_to_consider)):
            x1, y1, x2, y2 = boxes[j, :]
            color = get_color(j)
            score = scores[j]

            if args.display_bboxes:
                cv2.rectangle(img_numpy, (x1, y1), (x2, y2), color, 1)

            if args.display_text:
                _class = cfg.dataset.class_names[classes[j]]
                text_str = '%s: %.2f' % (_class, score) if args.display_scores else _class

                font_face = cv2.FONT_HERSHEY_DUPLEX
                font_scale = 0.6
                font_thickness = 1

                text_w, text_h = cv2.getTextSize(text_str, font_face, font_scale, font_thickness)[0]

                text_pt = (x1, y1 - 3)
                text_color = [255, 255, 255]

                cv2.rectangle(img_numpy, (x1, y1), (x1 + text_w, y1 - text_h - 4), color, -1)
                cv2.putText(img_numpy, text_str, text_pt, font_face, font_scale, text_color, font_thickness, cv2.LINE_AA)
            
    return img_numpy



def eval_onnx_with_nms(onnx_model_path):
    """
    使用导出的onnx带有nms的模型进行推理
    """
    print("\nRunning eval_onnx_with_nms\n")

    onnx_model = onnx.load(onnx_model_path)
    
    save_path = 'onnx.jpg' # 输出图片路径
    path = '4.jpg'        # 修改输入图片路径


    frame = torch.from_numpy(cv2.imread(path)).float()
    batch = FastBaseTransform()(frame.unsqueeze(0))
    
    # 检查模型
    try:
        onnx.checker.check_model(onnx_model)
        print("Model check passed.")
    except Exception as e:
        print(f"Model check failed: {e}")
    
    sess = rt.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])
    input_name = sess.get_inputs()[0].name
    loc_name = sess.get_outputs()[0].name
    conf_name = sess.get_outputs()[1].name
    mask_name = sess.get_outputs()[2].name
    priors_name = sess.get_outputs()[3].name
    proto_name = sess.get_outputs()[4].name
    proto_name2 = sess.get_outputs()[5].name

    with timer.env("ONNX Runtime"):
        preds = sess.run([loc_name, conf_name, mask_name, priors_name, proto_name, proto_name2], {input_name: batch.cpu().detach().numpy()})
    # preds是一个包含100*4 array的list
    # """
    # preds是一个列表, 包含以下元素:
    # boxes: (N, 4) --> 100, 4
    # mask: (N, 32) -->100, 32
    # class: (N,) --> 100
    # score: (N,) --> 100
    # proto: 138,138,32 --> 138,138,32 
    # priors: (4) --> 4
    # """

    preds_out = [{'box': torch.tensor(preds[0]), 
                  'mask': torch.tensor(preds[1]), 
                  'class': torch.tensor(preds[2]), 
                  'score': torch.tensor(preds[3]),
                  'proto': torch.tensor(preds[4]),
                  'priors': torch.tensor(preds[5])}]
    # for k,v in preds_out[0].items():
    #     print(k, v.shape)

    img_numpy = prep_display(preds_out, frame, None, None, undo_transform=False)
    
    if save_path is None:
        img_numpy = img_numpy[:, :, (2, 1, 0)]    
    cv2.imwrite(save_path, img_numpy)  # 保存图片
eval_onnx_with_nms('yolact.onnx')

最终推理结果:

相关推荐
丕羽42 分钟前
【Pytorch】基本语法
人工智能·pytorch·python
bryant_meng1 小时前
【python】Distribution
开发语言·python·分布函数·常用分布
m0_594526302 小时前
Python批量合并多个PDF
java·python·pdf
工业互联网专业3 小时前
Python毕业设计选题:基于Hadoop的租房数据分析系统的设计与实现
vue.js·hadoop·python·flask·毕业设计·源码·课程设计
钱钱钱端3 小时前
【压力测试】如何确定系统最大并发用户数?
自动化测试·软件测试·python·职场和发展·压力测试·postman
慕卿扬3 小时前
基于python的机器学习(二)—— 使用Scikit-learn库
笔记·python·学习·机器学习·scikit-learn
Json____3 小时前
python的安装环境Miniconda(Conda 命令管理依赖配置)
开发语言·python·conda·miniconda
小袁在上班3 小时前
Python 单元测试中的 Mocking 与 Stubbing:提高测试效率的关键技术
python·单元测试·log4j
白狐欧莱雅3 小时前
使用python中的pygame简单实现飞机大战游戏
经验分享·python·游戏·pygame
阿_旭3 小时前
基于YOLO11/v10/v8/v5深度学习的维修工具检测识别系统设计与实现【python源码+Pyqt5界面+数据集+训练代码】
人工智能·python·深度学习·qt·ai