经过调研大部分图像语义分割教程发现基本所有的教程内容都不够全面,有的教程只有搭建简单的网络实现语义分割训练功能;有的教程只有训练和预测部分;有的教程训练过程中只有损失函数没有精确度评价指标等等。鉴于此,本人经过调研多个图像语义分割教程并进行了详细的梳理总结,将图像语义分割所需要的知识点、数据集增强和训练技巧、多种评价指标、模型保存和转换等重要内容都进行了整合到该教程中,能够让初学者只需要学习这一个教程就能够完全掌握图像语义分割的基本功底。
一、训练图像分割网络主要流程
- 构建图像语义分割数据训练集、验证集和测试集
- 数据预处理、包括数据增强、数据标准化和归一化
- 从零构建图像语义分割网络模型
- 设置训练超参数,学习率、优化器、损失函数等超参数
- 模型训练、验证、预测、保存模型、模型onnx转化
二、各个流程简要说明
1. 构建图像语义分割数据训练集、验证集、测试集
在工程目录下,①新建datasets文件夹用于存放训练集、验证集和测试集;②在datasets文件夹内分别新建images和labels文件夹,用来放图片和对应的mask图片;③在images和labels文件夹内分别新建train、val、test文件夹用于存放图像语义分割的训练集、验证集和测试集对应的图像和mask标签,结构如下:

2.加载数据集,包括数据预处理、数据增强、数据标准化和归一化
3.图像语义分割网络模型从零构建
该语义分割网络模型采用了3层卷积下采样和3层上采样,并使用了conv1和conv2两层特征进行了两个简单的通道跳跃连接。网络模型结构如下图所示。

4.图像语义分割网络模型训练
训练过程中根据训练损失值保存当前最优的模型。


5.进行单张图片预测
首先加载保存好的训练模型到GPU,然后通过transforms制作单张图片测试数据格式作为模型预测输入,最后将模型预测输出进行mask判断和保存为图像格式。从模型预测的分割结果可以看出这个简单的图像语义分割模型对明显的人物像素区域可以分割正确,对背景一致的像素区域难以区分,需要对分割模型进行进一步的优化设计以实现更复杂像素的正确分割。

6.模型整体评价
首先制作需要评价模型的数据集,即提前准备好的测试集作为模型的整体评价。

7.模型onnx转换
将训练好的模型转换为onnx格式,以进行后续其他深度学习框架和平台的应用。ONNX(Open Neural Network Exchange,开放式神经网络交换格式)是一种模型文件格式,它在模型训练和模型推理中间提供了中间桥梁,使得上游不同的训练框架都能导出ONNX格式的模型,给到下游不同的推理框架都可以读取ONNX进行部署。

这种基于ONNX的模型训练,中间件,再到模型推理的方式使得:
- ONNX将模型训练和推理解耦,任意上游训练框架和下游推理框架都可以组合搭配,而不需要用同一种框架既进行训练又进行推理
- ONNX是通用的模型格式,不同训练框架输出的模型可以用ONNX作为桥梁进行转换,使得模型更方便迁移
- ONNX部署兼容性极强,支持多种推理框架,支持CPU/GPU推理,支持跨语言推理
- ONNX格式配合上类似ONNXRumtime等推理框架,相比于模型在原生环境的推理性能会有大幅的提升
onnx模型转换和模型转换后进行优化结果:

下图为转换前的last_model.pth和转换后的best_model.onnx模型结构对比:


详细参考链接: