Solov2-TensorRT-CPP
in this repo, we deployed SOLOv2 to TensorRT with C++.
See the video:https://www.bilibili.com/video/BV1rQ4y1m7mx
Requirements
- Ubuntu 16.04/18.04/20.04
- Cuda10.2
- Cudnn8
- TensorRT8.0.1
- OpenCV 3.4
- Libtorch 1.8.2
- CMake 3.20
Acknowledge
Getting Started
1. Install Solov2 from SOLO
download,and run it successfully
2. Export the ONNX model fron original model
**you can follow with ** SOLOv2.tensorRT.
That is, before export, you have to modify some parts of the original SOLOv2 first:
2.1. modify SOLO-master/mmdet/models/anchor_heads/solov2_head.py:154:0：
#Modify for onnx export, frozen the input size = 800x800, batch size = 1
size = {0: 100, 1: 100, 2: 50, 3: 25, 4: 25}
feat_h, feat_w = ins_kernel_feat.shape[-2], ins_kernel_feat.shape[-1]
feat_h, feat_w = int(feat_h.cpu().numpy() if isinstance(feat_h, torch.Tensor) else feat_h), int(feat_w.cpu().numpy() if isinstance(feat_w, torch.Tensor) else feat_w)
x_range = torch.linspace(-1, 1, feat_w, device=ins_kernel_feat.device)
y_range = torch.linspace(-1, 1, feat_h, device=ins_kernel_feat.device)
y, x = torch.meshgrid(y_range, x_range)
y = y.expand([1, 1, -1, -1])
x = x.expand([1, 1, -1, -1])
# Origin from SOLO
# x_range = torch.linspace(-1, 1, ins_feat.shape[-1], device=ins_feat.device)
# y_range = torch.linspace(-1, 1, ins_feat.shape[-2], device=ins_feat.device)
# y, x = torch.meshgrid(y_range, x_range)
# y = y.expand([ins_feat.shape[0], 1, -1, -1])
# x = x.expand([ins_feat.shape[0], 1, -1, -1])
2.2 single_stage_ins.py in the function of forward_dummy(), add the forward_dummy of mask, such as :
def forward_dummy(self, img):
x = self.extract_feat(img)
outs = self.bbox_head(x)
if self.with_mask_feat_head:
mask_feat_pred = self.mask_feat_head(
x[self.mask_feat_head.start_level:self.mask_feat_head.end_level + 1])
outs = (outs[0], outs[1], mask_feat_pred)
return outs
2.3 export onnx model move the onnx_exporter.py to the SOLO/demo/, then run
#kitti size
python onnx_exporter.py ../configs/solov2/solov2_light_448_r34_fpn_8gpu_3x.py ../weights/SOLOv2_light_R34.onnx --checkpoint ../checkpoints/SOLOv2_LIGHT_448_R34_3x.pth --shape 384 1152
3. build the tensorrt model
Firstly edit the config file:config.yaml
image_width: 1226
image_height: 370
# the position of exported ONNX model
onnx_path: "xxxx/SOLOv2_light_R34.onnx"
# save tensorrt model to:
serialize_path: "xxx/tensorrt_model_1152x384.bin"
# solo parameters
SOLO_NMS_PRE: 500
SOLO_MAX_PER_IMG: 100
SOLO_NMS_KERNEL: "gaussian"
#SOLO_NMS_SIGMA=2.0
SOLO_NMS_SIGMA: 2.0
SOLO_SCORE_THR: 0.1
SOLO_MASK_THR: 0.5
SOLO_UPDATE_THR: 0.2
segmentor_log_path: "xxx/log/segmentor_log.txt"
segmentor_log_level: "debug"
segmentor_log_flush: "debug"
# test img dir
DATASET_DIR: "xxx/kitti/odometry/colors/07/image_2/"
WARN_UP_IMAGE_PATH: "xxx/kitti.png"
and then,compile the CMake project:
mkdir build && cd build
cmake ..
make -j10
last, build the tensorrt model:
cd ..
./build/build_model ./config/config.yaml
4. run the demo
if you have the KITTI dataset, set config.yaml with right path DATASET_DIR ,run:
./build/InstanceSegment ./config/config.yaml
but if you not , and just want run at a image, set config.yaml with right image path WARN_UP_IMAGE_PATH, then run :
./build/demo ./config/config.yaml
![Parameter check failed at: runtime/api/executionContext.cpp::enqueueInternal::322, condition: bindings[x] != nullptr](https://avatars.githubusercontent.com/u/87312984?v=4)
58 Jan 3, 2023
70 Jan 4, 2023
46 Dec 6, 2022
120 Dec 18, 2022
2.4k Dec 30, 2022
675 Jan 6, 2023
172 Dec 29, 2022
16 Apr 19, 2022
34 Nov 25, 2022
10 Jul 30, 2021
35 Nov 10, 2022
1.5k Jan 5, 2023
48 Nov 25, 2022
7 Dec 10, 2022
43 Dec 27, 2022
62 Dec 14, 2022
2 Nov 7, 2021