本文代码见github:https://github.com/CamlinZ/Segmentation_Engineer
本文以https://github.com/CSAILVision/semantic-segmentation-pytorch作为示例语义分割算法,采用libtorch来进行C++部署
整个部署分为以下几个步骤进行展开:
- 示例语义分割算法简介
原始模型预处理opencv化
模型网络结构重构
- 模型转化
- C++数据预处理、前向传播及后处理代码开发
- opencv和libtorch联合编译
示例语义分割算法简介
https://github.com/CSAILVision/semantic-segmentation-pytorch整体工程代码经过简化整理后如下结构所示:
1 | Seg_Engeneer |
原始模型预处理opencv化
python版本的pytorch中针对图像任务通常采用torchvision来进行图像的预处理,而torchvision中的预处理函数是调用PIL库进行开发。由于python下的PIL库没有对应的C++版本,而opencv是同时有Python和C++版本的接口函数,所以这里为了保证模型的精度,需要将模型的预处理函数用opencv进行重构。
这里主要参考:https://github.com/jbohnslav/opencv_transforms
这里将torchvision中对应的PIL数据预处理操作全部转换为opencv,从接口函数的列表中可以看到是可以包含几乎所有的数据预处理操作的,同时接口的参数和torchvision基本一致。
1 | __all__ = [ |
本文示范的语义分割算法中,需要修改的代码主要涉及dataset.py文件,该文件中主要涉及到的是训练,验证和测试集的数据预处理操作,这里以训练数据的预处理操作进行示例:
这里需要将上述的opencv操作代码opencv_transforms按照第一步中展示的工程结构放入到工程中,将头文件改为
1 | from torchvision import transforms ---> from .opencv_transforms import transforms |
然后需要将resize函数修改为opencv的resize:
1 | def imresize(im, size, interp='bilinear'): |
最后需要对TrainDataset类中涉及到PIL的操作进行修改,这里主要涉及到__getitem__函数,经过修改后为:
1 | def __getitem__(self, index): |
模型网络结构重构
本次示例的语义分割网络中,作者为了方便选用不同的encoder和decoder网络结构,所以整体网络结构中的encoder和decoder被分成两个部分。在pth模型转换pt模型时,网络的输入和输出均要求是tensor的结构,而示例中encoder部分采用多尺度特征图给到decoder来做后续操作,所以encoder网络的输出为一个list结构,如下代码所示:
1 | def forward(self, x, return_feature_maps=False): |
所以这里需要将encoder和decoder的网络结构整合到一起,这里采用ResNet101作为encoder的网络结构,UperNet作为decoder的网络结构,合并网络后,即修改models文件夹下的models.py文件,修改后为:
1 | import torch |
以上图像预处理opencv化和网络模型结构修改完后,需要进行重新训练,得到修改后的模型
模型转化
本文中模型转换主要参考:https://blog.csdn.net/hiteryang/article/details/105575307中模型转化的思路,需要注意的是模型的输入和输出必须要是tensor,否则会出现转化失败的情况
C++数据预处理、前向传播及后处理代码开发
目前python训练的模型已经转换完毕,后面就是需要对照python前向转播的代码进行C++化,这里主要涉及到数据预处理、前向传播及后处理三个部分的代码
数据预处理
在第一节整体工程代码结构中,单张测试脚本test_img.py里的数据预处理代码为:
1 | def data_preprocess_opencv(img_path): |
1 | int round2nearest_multiple(int x, int p){ |
前向传播
在第一节整体工程代码结构中,单张测试脚本test_img.py里的前向传播代码为:
1 | img_ori = cv2.imread(img_path) |
转化后的C++代码为:
1 | cv::Mat input_image = cv::imread(img_path, cv::IMREAD_COLOR); |
opencv和libtorch联合编译
官网编译指引地址:https://pytorch.org/cppdocs/installing.html
按照以上编译流程进行操作即可,这里主要存在的问题是opencv和libtorch联合编译的问题,也即是Cmakelist.txt文件的写法,这里我opencv和libtorch都是安装在自定义的目录下:
1 | cmake_minimum_required(VERSION 3.0 FATAL_ERROR) |
最终编译通过之后就可以实现最终的结果了