2021全国数字生态创新大赛-智能算法赛
https://tianchi.aliyun.com/competition/entrance/531860
团队:wall墙
联系方式:18373180152 13431030112
1 运行环境
- 操作系统:Ubuntu 16.04.7 LTS(内核:Linux 4.15.0-129-generic x86_64)
- Python:3.7.3
- PyTorch:1.7.1
- CUDA:V10.2.89
- 需要安装的 Python package:详见
requirement.txt
2 解决方案
2.1 模型选择
我们团队前期对比了语义分割的如下几种
2.2 模型改进
2.2.1 backbone
- 实验中分别对多个网络 (见2.1) 使用不同的
$backbone$ $(efficientnet-b6$ 、$efficientnet-b7$、$resnet34$、$resnet50$、$resnext50-32x4d$、$timm{-}resnest101e$ 等$)$ 来做$encoder$ 提取图像特征信息 - 最后从
$mIoU$ 、模型参数及推理时间等指标权衡,选择$efficientnet-b7$ 作为$FPN$ 和$UnetPlusPlus$ 的$backbone$ ,$efficient-b6$ 作为$PSPNet$ 的$backbone$ 。
2.2.2 注意力机制
- 在分类/分割的时候,我们注意到遥感图像的一些特点,比如遥感图像多了个近红外波段,而使用近红外波段的 NDVI 指数能够很好区分植物,NDWI 指数能够区分水体。这里面又可能涉及一些问题,是不是都要用近红外波段来分类/分割,是不是有些类的分割只用 rgb 就能做很好,甚至有些单通道的区分信息就能胜任,加了其他通道反而淹没了区分信息。
- 于是我们采用了 Spatial-Channel Sequeeze & Excitation (SCSE) 是综合了通道维度和空间维度的注意力模块,对网络特征进行调整,自动学习到重要的特征图或者是特征通道对其加大权重,减小不重要特征的影响,从而改善图像分割的结果。
2.3 优化算法
- 我们使用了余弦退火暖重启的学习率调度算法配合
$AdamW$ 优化器,经过调参,验证准确率总是会在学习率的最低点达到一个很好的效果,而随着学习率回升,验证精度会有所下降.所以为了能最终得到一个更好的收敛点,设置$T_{mult}>1$ 是很有必要的,这样到了训练后期,学习率不会再有一个回升的过程,而且一直下降直到训练结束,最终设置$T_0$ 为$3$ ,$T_{mult}$ 为$2$ ,最低学习率为$1e-5$ 。
2.4 损失函数
-
$Cross Entropy Loss$ 可以用在大多数语义分割场景中,但它有一个明显的缺点,那就是当正负样本不平衡时,损失函数中数量多的样本的成分就会占据主导,使得模型在预测时严重偏向数量多的样本,导致效果不好。 - 为了削减正负样本不平衡的问题,我们加入了适用于语义分割任务的损失函数
$Dice Loss$ ,用来度量集合相似度的度量函数。 - 因为单独使用
$Dice Loss$ 作为损失函数可能会导致训练不稳定,故采用融合损失函数:
2.5 数据集处理
2.5.1 训练数据处理
- 当初可视化数据原图时,发现全通道图像比 rgb 图像要模糊,也考虑了将 (r, g, b, nir) 图片转为 (r, g, b) 图片来训练,但是效果不如使用全通道输入好,我们在复现过程中也曾加入通道注意力机制来解决这个问题。
2.5.2 通道归一化 / 均值滤波等
- 按照经验,对图像各通道做归一化处理能够加快模型收敛,但实验中发现结果变化不大,故舍弃。
- 训练过程中发现对道路、小河等一些细节区域分割精度较低,曾尝试基于
$Gamma$ 变换调整光照、基于均值滤波进行图像模糊化等处理,也因效果不是很好没有采用,代码详见code/utils/data_enhancement.py
。
2.5.3 数据增广
-
为了减少网络的过拟合现象,我们以
$50%$ 的概率对图片执行水平翻转、垂直翻转、旋转$90$ 度、转置这$4$ 个操作之一。又注意到需要对标签也做同样的操作,所以用了albumentations
库,代码如下:import albumentations as A transform = A.Compose([ A.Resize(IMAGE_SIZE, IMAGE_SIZE), A.OneOf([ A.VerticalFlip(p=0.5), A.HorizontalFlip(p=0.5), A.RandomRotate90(p=0.5), A.Transpose(p=0.5), ]) ])
-
训练过程中曾尝试以
$20%$ 的概率对图片进行$1%$ 的拉伸操作,但效果不是很好,故未采用,代码详见code/utils/image_stretch.py
2.6 模型训练
2.6.1 全数据训练
-
初始调试模型时将数据集按
$4: 1$ 的比例划分训练集与验证集进行模式训练。 -
中期为挑选表现更优模型,使用
$K-fold$ 交叉验证方法进行模型训练。 -
在整个训练流程确定下来后,我们使用了全数据进行重新训练,增大训练数据集,进一步增加模型的泛化能力。
2.6.2 随机权重平均
-
在训练过程中使用 Stochastic Weight Averaging (简称 SWA,随机权重平均) 来对训练过程中的多个checkpoints 进行平均,以寻找到更广的最优值域,提升模型的泛化能力。
-
经实验发现,SWA 在单模上起到了较好的表现效果,但拉低了模型集成的性能。
2.7 测试时增强
- 在测试时,为原始图像分别进行水平翻转和垂直翻转,并将它们(原图、原图经过水平翻转、原图经过垂直翻转)输入到模型中,然后对这多个版本的输出相加并计算平均结果作为图像的预测结果。
2.8 模型集成
- 我们使用以上训练**,训练了 $FPN、UNetPlusPlus $ 和 $ PSPNet$ 这
$3$ 个模型,单模在A榜的分数为$39.52 ~ 39.94$ 之间。最后,我们在预测时使用$bagging$ **进行集成,对每个模型的输出求取平均值作为最终结果,这极大可能的避免了不好的样本数据,从而提高了分割的准确度。
3 模型复现
3.1 训练流程
code/train.py
训练流程如下:
- 导入依赖包
- 设置随机种子(
seed = 6666
) - 定义计算
$mIoU$ 指标的函数 - 自定义
Dataset
类 - 自定义数据可视化
- 我们自定义了 10 个类别标签所映射的 rgb 颜色
- 定义可视化函数,以 原图 - 预测图 - 标签图 为一组可视化展示验证集结果
- 绘制
$loss$ 曲线以更好分析训练情况
- 配置训练参数
- epoch
- learning_rate
- optimizer
- criterion
- scheduler
- ...... (详见代码)
- 实例化
Dataset
并在其上运用数据增广- 水平翻转、垂直翻转、旋转90°、转置
- 开始训练
- 全数据训练,我们不执行 validation
- 保存 40 ~ 48 epoch 的模型参数
在项目主目录下运行如下代码即可复现模型训练
cd code
bash train.sh
3.2 预测流程
在项目主目录下运行如下代码即可复现模型预测
cd code
bash test.sh
4 额外说明
-
由于使用模型集成的策略,得到预测结果在 GeForce RTX 2080 Ti 下约要 45 mins。
-
训练模型复现过程中仅提供了
$UNetPlusPlus$ 的训练过程,若要继续复现$FPN$ 、$PSPNet$ 的训练过程,可在code/train.py
文件中的Model parameters
代码段 (第167 ~ 172
行) 将注释的代码取消注释。 -
只保存了
$FPN$ 和$UNetPlusPlus$ 的模型文件。 -
保存的模型参数是在 cuda 上进行训练,在训练中需要指定 cuda 则进入
code/train.py
修改第164
行代码,预测中需要指定 cuda 则进入code/test.py
修改第29
行代码。 -
比赛中的代码规范示例
tcdata
中测试集所在路径为tcdata/suichang_round2_test_partB_210316/
我们下载的测试集文件夹名称为
suichang_round1_test_partB_210120
,如果需要修改则进入code/test.py
修改第60
行。