124 lines
5.7 KiB
Markdown
124 lines
5.7 KiB
Markdown
## SegFormer语义分割模型在Pytorch当中的实现
|
||
---
|
||
|
||
### 目录
|
||
1. [仓库更新 Top News](#仓库更新)
|
||
2. [相关仓库 Related code](#相关仓库)
|
||
3. [性能情况 Performance](#性能情况)
|
||
4. [所需环境 Environment](#所需环境)
|
||
5. [文件下载 Download](#文件下载)
|
||
6. [训练步骤 How2train](#训练步骤)
|
||
7. [预测步骤 How2predict](#预测步骤)
|
||
8. [评估步骤 miou](#评估步骤)
|
||
9. [参考资料 Reference](#Reference)
|
||
|
||
## Top News
|
||
**`2022-06`**:**创建仓库、支持训练时评估、支持多backbone、支持step、cos学习率下降法、支持adam、sgd优化器选择、支持学习率根据batch_size自适应调整。**
|
||
|
||
## 相关仓库
|
||
| 模型 | 路径 |
|
||
| :----- | :----- |
|
||
Unet | https://github.com/bubbliiiing/unet-pytorch
|
||
PSPnet | https://github.com/bubbliiiing/pspnet-pytorch
|
||
deeplabv3+ | https://github.com/bubbliiiing/deeplabv3-plus-pytorch
|
||
hrnet | https://github.com/bubbliiiing/hrnet-pytorch
|
||
segformer | https://github.com/bubbliiiing/segformer-pytorch
|
||
|
||
### 性能情况
|
||
| 训练数据集 | 权值文件名称 | 测试数据集 | 输入图片大小 | mIOU |
|
||
| :-----: | :-----: | :------: | :------: | :------: |
|
||
| VOC12+SBD | [segformer_b0_weights_voc.pth](https://github.com/bubbliiiing/segformer-pytorch/releases/download/v1.0/segformer_b0_weights_voc.pth) | VOC-Val12 | 512x512 | 73.34 |
|
||
| VOC12+SBD | [segformer_b1_weights_voc.pth](https://github.com/bubbliiiing/segformer-pytorch/releases/download/v1.0/segformer_b1_weights_voc.pth) | VOC-Val12 | 512x512 | 76.80 |
|
||
| VOC12+SBD | [segformer_b2_weights_voc.pth](https://github.com/bubbliiiing/segformer-pytorch/releases/download/v1.0/segformer_b2_weights_voc.pth) | VOC-Val12 | 512x512 | 80.38 |
|
||
|
||
### 所需环境
|
||
torch==1.2.0
|
||
|
||
### 文件下载
|
||
训练所需的权值可在百度网盘中下载。
|
||
链接: https://pan.baidu.com/s/1tH4wdGnACtIuGOoXb0_rAw
|
||
提取码: tyjr
|
||
|
||
VOC拓展数据集的百度网盘如下：
|
||
链接: https://pan.baidu.com/s/1vkk3lMheUm6IjTXznlg7Ng
|
||
提取码: 44mk
|
||
|
||
### 训练步骤
|
||
#### a、训练voc数据集
|
||
1、将我提供的voc数据集放入VOCdevkit中（无需运行voc_annotation.py）。
|
||
2、在train.py中设置对应参数，默认参数已经对应voc数据集所需要的参数了，所以只要修改backbone和model_path即可。
|
||
3、运行train.py进行训练。
|
||
|
||
#### b、训练自己的数据集
|
||
1、本文使用VOC格式进行训练。
|
||
2、训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的SegmentationClass中。
|
||
3、训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
|
||
4、在训练前利用voc_annotation.py文件生成对应的txt。
|
||
5、在train.py文件夹下面，选择自己要使用的主干模型。
|
||
6、注意修改train.py的num_classes为分类个数+1。
|
||
7、运行train.py即可开始训练。
|
||
|
||
### 预测步骤
|
||
#### a、使用预训练权重
|
||
1、下载完库后解压，在百度网盘下载权值，放入model_data，修改segformer.py的backbone和model_path之后再运行predict.py，输入。
|
||
```python
|
||
img/street.jpg
|
||
```
|
||
可完成预测。
|
||
2、在predict.py里面进行设置可以进行fps测试、整个文件夹的测试和video视频检测。
|
||
|
||
#### b、使用自己训练的权重
|
||
1、按照训练步骤训练。
|
||
2、在segformer.py文件里面，在如下部分修改model_path、num_classes、backbone使其对应训练好的文件；**model_path对应logs文件夹下面的权值文件，num_classes代表要预测的类的数量加1，backbone是所使用的主干特征提取网络**。
|
||
```python
|
||
_defaults = {
|
||
#-------------------------------------------------------------------#
|
||
# model_path指向logs文件夹下的权值文件
|
||
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
|
||
# 验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。
|
||
#-------------------------------------------------------------------#
|
||
"model_path" : "model_data/segformer_b0_weights_voc.pth",
|
||
#----------------------------------------#
|
||
# 所需要区分的类的个数+1
|
||
#----------------------------------------#
|
||
"num_classes" : 21,
|
||
#----------------------------------------#
|
||
# 所使用的的主干网络：
|
||
# b0、b1、b2、b3、b4、b5
|
||
#----------------------------------------#
|
||
"phi" : "b0",
|
||
#----------------------------------------#
|
||
# 输入图片的大小
|
||
#----------------------------------------#
|
||
"input_shape" : [512, 512],
|
||
#-------------------------------------------------#
|
||
# mix_type参数用于控制检测结果的可视化方式
|
||
#
|
||
# mix_type = 0的时候代表原图与生成的图进行混合
|
||
# mix_type = 1的时候代表仅保留生成的图
|
||
# mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
|
||
#-------------------------------------------------#
|
||
"mix_type" : 0,
|
||
#-------------------------------#
|
||
# 是否使用Cuda
|
||
# 没有GPU可以设置成False
|
||
#-------------------------------#
|
||
"cuda" : True,
|
||
}
|
||
```
|
||
3、运行predict.py，输入
|
||
```python
|
||
img/street.jpg
|
||
```
|
||
可完成预测。
|
||
4、在predict.py里面进行设置可以进行fps测试、整个文件夹的测试和video视频检测。
|
||
|
||
### 评估步骤
|
||
1、设置get_miou.py里面的num_classes为预测的类的数量加1。
|
||
2、设置get_miou.py里面的name_classes为需要去区分的类别。
|
||
3、运行get_miou.py即可获得miou大小。
|
||
|
||
### Reference
|
||
https://github.com/ggyyzm/pytorch_segmentation
|
||
https://github.com/bonlime/keras-deeplab-v3-plus
|
||
https://github.com/NVlabs/SegFormer |