yolact导出onnx

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

一、导出的代码

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

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

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。

use_jit = torch.cuda.device_count() <= 1
use_jit = False   
if not use_jit:
    print('Multiple GPUs detected! Turning off JIT.')

3. numpy

RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.

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

# 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：

self.use_fast_nms = True
# self.use_fast_nms = False

4. tupels

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。

# out.append({'detection': result, 'net': net})
out.append(result)

output_utils.py中注释掉net

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]

5. output_utils.py

第74行增加这一句，输出增加 ‘priors’，不然后续推理出错：

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，不会执行。只是为了后面我的验证代码能够正常运行：

# 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)

6. box_utils.py
   (很重要）这里要将@torch.jit.script注释掉，否则到处的结果是错误的：

# @torch.jit.script
def decode(loc, priors, use_yolo_regressors:bool=False):

至此可以导出。

三、验证

注意修改图片路径：

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')

最终推理结果：