如何用自己的数据训练MASK R-CNN模型

AiTechYun

编辑：chux

计算机一直很擅长数字处理，却苦于分析图像中的大量数据。直到最近，创建图形处理单元库已不仅仅用于游戏，现在我们可以利用数千个核心的原始力量来揭示图片背后的含义。

使用你的数据

我们将以形状数据集作为范例，其中颜色和大小随机的圆形、正方形和三角形分布在颜色随机的背景上。我们之前已经创建了一个COCO类型的数据集。如果你想学习如何转换自己的数据集，请查看如何用pycococreator将自己的数据集转换为COCO类型。

这次的重点将是自动标记图像中的所有形状，并找出每个图形的位置，精确到像素。这种任务叫做“对象分割”。在探索计算机视觉过程中，你可能也遇到过诸如“对象识别”、“类分割”和“对象检测”之类的术语。这些术语听起来很相似，一开始你可能会感到迷惑，那么就让我们看看它们各自的功能吧。下面是我们从每种类型中获取的信息示例，任务难度从左到右递增。

对象识别告诉我们图像中有什么，但不会提示位置和数量。类分割将位置信息添加到图像的不同类型的对象中。对象检测则将每个对象用边界框分隔开。这四种当中难度最大的就是我们要进行训练的对象分割。它为每个形状提供清晰的边界，这也能够得到前三种功能的结果。

使用这个简单的数据集，我们可以使用旧的计算机视觉概念，如用于检测圆和线的Hough或模板匹配来得到合适的结果。但使用深度学习，我们不需要改变方法，就可以在任何类型的图像数据集中得到相同类型的结果，而且都不用考虑我们要找的具体特征，这就像魔法一般神奇。

什么是MASK R-CNN？

在我们开始训练自己的Mask R-CNN模型前，首先来搞清楚这个名称的含义。我们从右到左来介绍。

“NN”就是指神经网络，这一概念受到了对生物神经元是如何工作的想象的启发。神经网络是连接到一起的神经元的容器，每个神经元根据其输入和内部参数输出信号。当我们训练神经网络时，我们调整神经元的内部参数，以便得到符合期望的输出。

“C”代表“卷积的”，CNN是专门为用图像学习而设计的，但与标准的神经网络相似。它学习的过滤器每次都以一小块为单位上下滑动（卷积）图像，而不是一次浏览整个图像。与常规的神经网络相比，CNN需要的参数和内存占用更少，这使得它们可以超越传统神经网络的限制，处理尺寸更大的图像。

简单的CNN擅长对象识别，但如果我们想要进行对象检测，我们需要知道其位置。所以R，就是区域位置。R-CNN能够在找到的物体周围绘制出边界框。现在的R-CNN已经发展得比原始的R-CNN更快，因此它们被称为“Fast R-CNN”和“Faster R-CNN”。Faster R-CNN在CNN终端加入了Region Proposal Network用来计划区域。如果在区域中发现了一个对象，那么就将这些区域作为边界框。

最后来解释一下“Mask”，它增加了像素级的分割，并创建了对象分割模型。它在网络中添加了一个额外的分支以创建二进制掩码，这与我们注释图像的做法类似。

介绍完了这一概念不同部分的含义和作用，现在让我们开始训练自己的Mask-RCNN。

准备好电脑

要运行这些示例，你需要Ubuntu 16.04系统和近期的nvidia显卡。我使用只有2 GB内存的GeForce 940 M来训练一小部分的网络，推荐你用11 GB或更多内存的nvidia显卡。如果你没有这些条件，你可以从Amazon Web Services或Google Cloud开始。

为了确保同步，我们将使用Docker来运行所有示例。Docker使用脚本创建系统的副本，因此你无需担心自己安装的问题。但是，在我们能够因自动设置变得轻松之前，首先我们需要准备好主机系统。这部分可能有点麻烦，但绝对值得。

装好Ubuntu 16.04系统后，我们需要安装nvidia显卡驱动以及CUDA (并行计算平台)。首先打开一个终端，运行以下命令来安装显卡驱动。

然后安装CUDA。

wget https://developer.nvidia.com/compute/cuda/9.1/Prod/local_installers/cuda-repo-ubuntu1604-9-1-local_9.1.85-1_amd64

sudo dpkg -i cuda-repo-ubuntu1604-9-1-local_9.1.85-1_amd64.deb

sudo apt-key add /var/cuda-repo-<version>/7fa2af80.pub

sudo apt-get update

sudo apt-get install cuda-toolkit-9-1 cuda-libraries-dev-9-1 cuda-libraries-9-1

现在来安装Docker，Docker-Compose和Nvidia-Docker。

wget -O get_docker.sh get.docker.com

chmod +x get_docker.sh

sudo ./get_docker

sudo apt-get install -y docker-ce=18.03.0~ce-0~ubuntu

sudo usermod -aG docker $USER

sudo curl -L https://github.com/docker/compose/releases/download/1.21.0/docker-compose-$(uname -s)-$(uname -m) -o /usr/local/bin/docker-compose

sudo chmod +x /usr/local/bin/docker-compose

curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | \ sudo apt-key add -

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)

curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | \ sudo tee /etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update

sudo apt-get install nvidia-docker2

sudo pkill -SIGHUP dockerd

需要运行的命令很多，希望你那边能够顺利安装。现在我们终于可以运行Mask R-CNN系统了，重启电脑确保一切能够正常运行。

开始探索

在这个网站(github.com/waspinator/deep-learning-explorer)下载并解压deep-learning-explorer。你会找到mask-rcnn文件夹和一个数据文件夹。另一个压缩文件中有我们的测试数据集。解压shapes.zip文件，将annotations，shapes_train2018，shapes_test2018和shapes_validate2018移到data/shapes中。

回到终端，cd进入mask-rcnn/docker，然后运行dockdocker。首次运行这个命令时，DockerDocker将从头开始创建系统，因此可能需要等几分钟，这之后就很快了。一旦准备好了你就会看到这样的提示：

把最后一行粘贴到浏览器中，你就会进入Jupyter Notebook网站。依次点击home/keras/mask-rcnn/notebooks/mask_rcnn.ipynb，现在你可以遍历所有Notebook，并训练自己的Mask R-CNN模型。用Keras和Tensorflow在后台训练gpu上的神经网络。如果你没有11GB内存的显卡，那么你可以在调整这一步上出现问题，不过就算是只有2GB内存的显卡，也可以训练网络最上面的部分。

我们不用花费数天或数周的时间来训练模型，也没有成千上万的例子，但我们还能得到相当好的结果，是因为我们从真正的COCO数据集之前的训练中复制了权重（内部神经元参数）。由于大多数图像数据集都有相似的基本特征，比如颜色和模式，所以训练一个模型得出的数据通常可以用来训练另一个模型。以这种方式复制数据的方法叫做迁移学习。

在notebook的底部你会发现预测出正确形状的概率只有37%。你可以将STEPS_PER_EPOCH增加到750（训练样本总数），并且运行5次或更多以使模型做出更好的预测。

在训练期间或训练后，你可以用TensorBoard查看图表确认进展。我们需要登录我们刚刚启动的Docker容器，然后在web浏览器中访问它之前运行TensorBoard。在终端运行docker ps，这样你就能看到所有运行中的容器。使用CONTAINER ID的前两个字符启动训练模型的Docker容器中的bash shell。例如，ID是d5242f7ab1e3，我们使用docker exec -it d5 bash登录。现在我们已经在容器中了，运行tensorboard –logdir ~/data/shapes/logs –host 0.0.0.0，这样你就可以访问http://localhost:8877获取TensorBoard。

现在尝试一下用自己的数据来训练Mask R-CNN模型吧。