Skip to content

Latest commit

 

History

History
103 lines (71 loc) · 4.01 KB

README.md

File metadata and controls

103 lines (71 loc) · 4.01 KB
Raw

Mask RCNN for TensorRT

这是一个将Keras模型转化为.uff的通用的一种做法！过程中使用了TensorRT的plugin,为了使用plugin,将Keras模型转换为Tensorflow的.pb模型，然后将.pb模型通过GraphSurgeon 和UFF转换为.uff模型.

主要涉及的Plugin

  • ResizeNearest - Nearest neighbor interpolation for resizing features. This works for the FPN (Feature Pyramid Network) module.
  • ProposalLayer - Generate the first stage's proposals based on anchors and RPN's (Region Proposal Network) outputs (scores, bbox_deltas).
  • PyramidROIAlign - Crop and resize the feature of ROIs (first stage's proposals) from the corresponding feature layer.
  • DetectionLayer - Refine the first stage's proposals to produce final detections.
  • SpecialSlice - A workaround plugin to slice detection output [y1, x1, y2, x2, class_id, score] to [y1, x1, y2 , x2] for data with more than one index dimensions (batch_idx, proposal_idx, detections(y1, x1, y2, x2)).

注意：我们参考了TensorRT官方sample中的Mask RCNN的实现，请按照本Repo中要求配置环境，这样避免出现各种各样的错误（菜坑后的忠告），其次我们提供了独立易读的C++代码调用，不会像TensorRT sample中因为多层封装不易阅读。

1.模型说明

2.环境搭建

  • Docker镜像： nvcr.io/nvidia/tensorflow:18.12-py3 # tensorflow=1.12
  • uff=0.6.5
  • graphsurgeon=0.4.1
  • imageio=2.8.0
  • keras == 2.1.3
  • scikit-image

3.安装环境及TensorRT序列化

  1. 安装必要的library

    我们基于基础的Tensorflow镜像进行安装nvcr.io/nvidia/tensorflow:18.12-py3

sudo nvidia-docker run  nvcr.io/nvidia/tensorflow:18.12-py3 
cd ./converted/
pip3 install -r requirements.txt
pip3 install ./pylib/uff-0.6.5-py2.py3-none-any.whl
pip3 install ./pylib/graphsurgeon-0.4.1-py2.py3-none-any.whl
  1. 修改uff库的代码：一般在 /usr/lib/python3.6/dist-packages/uff/converters/tensorflow/converter_functions.py.
uff_graph.conv_transpose(
    inputs[0], inputs[2], inputs[1],
    strides, padding,
    dilation=None, number_groups=number_groups,
    left_format=lhs_fmt, right_format=rhs_fmt,
    name=name, fields=fields
    )
  1. 下载keras版Mask RCNN的代码库
git clone https://github.com/matterport/Mask_RCNN.git
  1. 使用patch更新Mask RCNN代码库的代码，将代码中的NHWC转为NCHW
cd Mask_RCNN
git checkout 3deaec5
patch -p1 < ../0001-Update-the-Mask_RCNN-model-from-NHWC-to-NCHW.patch
cd -
  1. 下载预训练的模型
wget https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5
  1. 将Keras模型转.pb并转为.uff
# 自己的模型得修改config.py!!!!
python3 mrcnn_to_trt_single.py -w ../mask_rcnn_coco.h5 -o ../mrcnn_nchw.uff -p ./config.py
  1. trtexec 生成engine
trtexec --uff=mrcnn_nchw.uff --output=mrcnn_detection --output=mrcnn_mask/Sigmoid --uffInput=input_image,3,1024,1024 --saveEngine=mask.plan --workspace=3000 --verbose

看到了pass恭喜你成功序列化Mask RCNN.

5.C++调用TensorRT模型进行识别

我们提供了Mask RCNN在C++下通过TensorRT调用识别的完整代码，其识别结果如下图所示: