yolact导出onnx

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

一、导出的代码

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

import torch
import cv2from yolact import Yolactdef 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 timerdef 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.0h, w, _ = img.shape# 后处理 w, h = 612 612with 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 = jbreak# Quick and dirty lambda for selecting the color for a particular index# Also keeps track of a per-gpu color cache for maximum speeddef get_color(j, on_gpu=None):global color_cachecolor_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, dependingcolor = (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] = colorreturn 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 sufficeif 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 isinv_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_cumulmasks_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 reasonimg_numpy = (img_gpu * 255).byte().cpu().numpy()if num_dets_to_consider == 0:return img_numpyif 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 _classfont_face = cv2.FONT_HERSHEY_DUPLEXfont_scale = 0.6font_thickness = 1text_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_numpydef 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].nameloc_name = sess.get_outputs()[0].nameconf_name = sess.get_outputs()[1].namemask_name = sess.get_outputs()[2].namepriors_name = sess.get_outputs()[3].nameproto_name = sess.get_outputs()[4].nameproto_name2 = sess.get_outputs()[5].namewith 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')

最终推理结果：