轻松学Pytorch之Faster-RCNN模型ONNX导出与部署

点击上方蓝字关注我们

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

前言

自从我写了这个系列文章以后，已经快两年时间了，我经常被人问到pytorch中的Faster-RCNN是否可以导出ONNX格式，这个问题后来pytorch官方有文档了，可以直接导出。后来第二个问题就是导出ONNX格式文件无法部署，其实原因在于第一条是因为官方导出的那个模式是个超像素的，直接把脚本拿过来用是能导出Faster-RCNN但是无法被ONNXRUNTIME使用，导致后来一系列的问题都无解了。这个就是所谓的死读书害死人！所以我自己尝试了一下，发现是可以导出ONNX格式的，直接使用torchvision提供的预训练Faster-RCNN模型，通过脚本就可以直接导出ONNX格式模型，而且还能部署。然后我又测试了一下自定义的模型，发现也没问题。因此特分享一下给大家！

Faster-RCNN转ONNX

这一步我主要使用pytorch自带的torch.onnx.export方法，该函数常见相关参数的说明如下：

代码语言：javascript复制

model, // 模型文件args, // 输入图像f, // 保存模型文件export_params=True, // 导出全部参数verbose=False, // 默认Falsetraining=TrainingMode.EVAL, // 推理模型input_names=None, // 输入节点名称output_names=None, // 输出节点名称opset_version=None // op版本do_constant_folding=True //dynamic_axes=None // 是否支持动态大小输入

把Faster-RCNN转行为ONNX模型的脚本如下：

代码语言：javascript复制

model = tv.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) dummy_input = torch.randn(1, 3, 1333, 800) model.eval() model(dummy_input) im = torch.zeros(1, 3, 1333, 800).to("cpu") torch.onnx.export(model, im,                     "faster_rcnn.onnx",                     verbose=False,                     opset_version=11,                     training=torch.onnx.TrainingMode.EVAL,                     do_constant_folding=True,                     input_names=['input'],                     output_names=['output'],                     dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}})

其中最后一个dynamic_axes参数，表表示第0、第2、第3支持任意尺度的输入，意思就是支持任意大小图像输入。

ONNXRUNTIME部署运行

代码贴出来了，自己看吧，很多函数都是经常使用的，看着就眼熟，代码就是先创建ONNXRUNTIME对象，然后加载模型，加载图像预处理为NCHW格式，然后转换为0~1之间，执行推理，后处理输出结果并显示。

代码语言：javascript复制

transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

sess_options = ort.SessionOptions()
# Below is for optimizing performance
sess_options.intra_op_num_threads = 24
# sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_session = ort.InferenceSession("faster_rcnn.onnx", sess_options=sess_options)
src = cv.imread("D:/images/cars.jpg")
image = cv.cvtColor(src, cv.COLOR_BGR2RGB)
blob = transform(image)
c, h, w = blob.shape
input_x = blob.view(1, c, h, w)
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_x)}
ort_outs = ort_session.run(None, ort_inputs)
boxes = ort_outs[0] # boxes
labels = ort_outs[1] # labels
scores = ort_outs[2] # scores
print(boxes.shape, boxes.dtype, labels.shape, labels.dtype, scores.shape, scores.dtype)

index = 0
for x1, y1, x2, y2 in boxes:
    if scores[index] > 0.5:
        cv.rectangle(src, (np.int32(x1), np.int32(y1)),
                     (np.int32(x2), np.int32(y2)), (0, 255, 255), 1, 8, 0)
        label_id = labels[index]
        label_txt = coco_names[str(label_id)]
        cv.putText(src, label_txt, (np.int32(x1), np.int32(y1)), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)
    index  = 1
cv.imshow("Faster-RCNN Detection Demo", src)
cv.waitKey(0)
cv.destroyAllWindows()

运行结果如下：