基于MindSpore框架的室内场景图像分割方法研究

概述

本文以华为最新国产深度学习框架Mindspore为基础，研究室内场景语义分割方法。本文基于注意力机制改进U-Net网络，并选取VGG16与ResNet50作为骨干网络，并且利用VGG16与ResNet50的预训练权重进行迁移学习。整体的技术路线如图1所示。

图片1

项目地址

https://gitee.com/zou-tianyu/mind-spore-unet-tianyu.git

名称	配置信息
NPU	Ascend910
操作系统	Ubuntu 20.04
编译器	Python3.7
框架	MindSpore1.6

NYU-V2数据集

本文利用公开数据集NYU-V2作为实验数据集，选取该数据集被标注的RGB图片用于训练和测试。作为常用的语义分割数据集，NYU-V2面向各种类型任务，选取的数据是由微软的 RGB摄像机记录的各种室内场景图片组成，其中共有464个不同的室内场景、1449张图片、894个类别标签。由于在NYU-V2数据集中，接近70%的数据集标签由前10类标签组成，所以本文选取该数据集室内场景的主要语义类别：墙壁（wall）、地板（floor）、橱柜（cabinet）、床（bed）、椅子（chair）、沙发（sofa）、桌子（table）、门（door）、窗户（window）、书柜（bookshelf）作为训练分割的语义类别，NYU-V2数据集中其他类别归为背景（background）类

数据集下载地址：NYU Depth V2 « Nathan SilbermanNYU Depth V2 « Nathan Silberman

实现流程记录

1、基于U-Net网络实现室内场景语义分割模型

U-Net是由Ronneberger 等人在2015年构建的一种完全对称的U型结构神经网络。参考华为官方gitee仓库的modelzoo克隆至本地，找到research/cv/unet目录，在此代码基础上进行修改。在原有U-Net模型的基础上，本文将输入图片的大小统一归为512×512，便于网络的卷积计算。本文实现的U-Net网络结构如图2所示。

2、模型权重文件转换

MindSpore采用ModelCheckpoint对象保存模型参数到CheckPoint文件（简称ckpt文件），由于MindSpore目前没有VGG16和ResNet50的相关预训练模型，导致本文在迁移学习的过程中无法直接使用MindSpore官方提供的预训练权重。PyTorch上采用pth文件保存模型参数且与MindSpore在保存模型参数的机制上不同，无法直接采用PyTorch上的预训练权重加载到MindSpore中。

本文在查阅了MindSpore与PyTorch相关文档后，成功编写了将pth模型文件转换为ckpt模型文件的脚本。经过实际训练发现该脚本切实有效的解决了MindSpore中预训练模型缺失的问题。

3、优化器的选择

优化器运用在神经网络的反向传播中，通过引导损失函数的参数向正确的方向更新适当的大小，使得损失函数的值能够到达全局最优。其中最为重要的两个部分是优化方向和步长。本文选取Adam作为模型训练的优化器，Adam是SGD、AdsGrad等优化器的集大成者，能够有效控制学习率步长和梯度方向。

本文由于运用到了迁移学习故将整个训练分为两个阶段。第一阶段冻结骨干网络预训练权重，将主要资源运用在扩张路径上，从而缩短训练时间，节约训练资源。第二阶段解冻骨干网络的参数，让网络中的所有参数参与训练。通过冻结训练后再解冻可以加快训练效率同时防止骨干网络的参数权值被破坏，从而获得更好的训练效果。

结果展示