Unet语义分割训练和TensorRT部署

参考链接：https://gitee.com/dogeblog/unet#%E4%BB%8B%E7%BB%8D
防止失效，以下记录：
介绍
基于unet的语义分割。 包括标注，标注结果转训练集，数据集训练，训练结束后转tensorrt,tensorrt客户端部署。

tensorrt部署后，640*360的图，在rtx3090推理时间4-5ms

原始代码来源
unet代码源于https://github.com/milesial/Pytorch-UNet
tensorrt代码源于https://github.com/wang-xinyu/tensorrtx 与 https://github.com/YuzhouPeng/unet-tensorrt
1.标注
参考https://gitee.com/doge_ac_cn/labelme#installation

2.生成unet训练的数据集
1.先把labelme标注的结果放在 labelme绘制mask/labelme 文件夹下。

2.然后先后运行labelme2coco.py和genneratemask.py，分别生成coco标注文件和mask文件。

3.训练
1.把原图和mask分别拷贝到 Pytorch-UNet/data 的img和mask文件夹下

2.运行train.py

3.命令行输入tensorboard --logdir=runs

然后进入http://localhost:6006/ 即可查看运行的loss值，学习率，验证集效果。
4.转tensorrt 的 推理 engine
看unet2tensorrt的readme即可

这里直接转的话，会与python的结果有误差

详情见作者的issue https://github.com/wang-xinyu/tensorrtx/issues/526 与 https://github.com/YuzhouPeng/unet-tensorrt/issues/8

作者代码链接：
https://download.csdn.net/download/zhouzongzong/86400386

5、Cmake文件（关于如何用Cmake配置opencv C++环境）
参考链接：https://blog.csdn.net/Wuhao9_/article/details/123980679

cmake_minimum_required(VERSION 2.6)

project(unet)

add_definitions(-std=c++11)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)

# cuda directory
include_directories(${PROJECT_SOURCE_DIR}/include)
include_directories(C:/Program Files/NVIDIA GPU\ Computing\ Toolkit/CUDA/v11.4/include)
link_directories(C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.4/lib/x64)

# tensorrt
include_directories(F:/yolov5_tensorRT/TensorRT-8.0.3.4/include)
link_directories(F:/yolov5_tensorRT/TensorRT-8.0.3.4/lib)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")

# link library and add exec file
add_executable(unet ${PROJECT_SOURCE_DIR}/unet.cpp)
target_link_libraries(unet nvinfer)
target_link_libraries(unet cudart)

add_definitions(-O2 -pthread)

# opencv library
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})
target_link_libraries(unet ${OpenCV_LIBS})

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

Cmake后生成文件：

6、错误定位：基本运行检查改为默认值（D8016:O2和RTC1命令行选项不兼容）
下面命令行删除，否则出现，error D8021: 无效的数值参数

7、运行测试：
运行前需要更改路径权重参数路径：

确认使用推理的精度：

利用unet -s命令生成引擎文件unet.engine

利用unet -d …/samples 进行测试，kTF32精度时间
利用unet -d …/samples 进行测试， kFP16精度时间