本文入选技术写作营优秀文章,作者李剑锋
MMDetection 是一个由 OpenMMLab 开发的开源目标检测工具箱,基于 PyTorch 实现。该库提供了丰富的目标检测算法,包括经典的 Faster R-CNN、YOLO 和最新的一些研究成果,非常方便于研究者和工程师进行模型的训练和推理。具有高度模块化和可扩展性的设计,使得用户可以非常灵活地进行个性化配置和二次开发。这一工具箱已经成为目标检测领域的事实标准之一,被广泛应用于学术研究和产业界。
(文末点击阅读原文可直达,欢迎体验)
下面,我将一步步从环境配置和数据采集开始介绍如何训练一个自己的 MMDetection 模型。
模型配置及训练
配置环境一直以来都是拦在新手面前的最大障碍,幸运的是 MMDetection 文档给出了详细的安装教程,我们跟随着文档的脚步一步步进行安装即可。
首先我们需要利用 conda 代码创建一个名为 mmdetection 的虚拟环境,并利用代码激活环境。
由于 OpenMMLab 算法是基于 PyTorch 进行开发的,因此我们需要安装 PyTorch 以实现其功能。根据官方文档,我们可以通过以下方式安装 PyTorch。一般来说,因为我们已经安装好 CUDA 和 Cudnn,因此我们选择安装 GPU 版本即可。
我们可以通过 conda 自带的 python 来测试是否成功安装 PyTorch 的 GPU 版本。如图所示,我在终端输入 python 后,进入了 Python 的环境并输入以下代码。假如出现了 False,则代表安装失败。假如是 True 则代表安装成功。我这里显示的是 False,就代表我安装失败了。因此我就需要卸载掉这个版本的 PyTorch 然后再在官网找到对应的版本重新安装。最后我们还需要输入 exit() 来退出 Python 环境回到终端。
卸载 PyTorch 的指令如下:
卸载完后,我们需要前往 PyTorch的官网 找到合适的版本进行下载。可以看到,PyTorch官方提供了很多种方式让我们选择,包括 CUDA 的版本,OS 系统以及安装的方式和语言。由于目前只支持到 CUDA11.8,高于 CUDA11.8 的都下载该版本,因此我们需要选择该方式进行下载(对于之前版本的 CUDA,可以点击框选部分的链接进行查看)。
安装完成后,我们再来测试一下是否成功安装。
可以看到,现在显示的就是 True 了。这也就意味着 PyTorch 的 GPU 版本成功被安装了。我们就可以继续下一步的安装操作了。
安装完成后,我们还需要下载 GitHub 的源码到本地并安装对应的配置文件。但是在此之前,假如 C 盘的空间紧张,我建议还是更换地址后再进行下载。
假如由于网络问题无法成功下载,我们可以去 MMDetection GitHub,并找到 ZIP 文件进行下载。然后把文件放到我们想要放的位置,然后再 cd到 该文件夹里,最后执行 pip install -v -e . 操作即可。
假如安装成功后将会显示以下界面:
在安装完后,官网还提供一个代码来测试是否成功安装。
假如上述的操作成功被执行而且没有报错,那么恭喜你,你成功安装了 MMDetection!接下来就可以正式进入开发的工作了。
数据准备
在本项目中,我们将使用开源的数据集来演示 MMDetection 的全流程使用。常见的开源数据集网站包括 Hugging Face、Kaggle 和 AI Studio。由于我是土木工程管理专业的,在该专业里最常见的目标检测应用就是安全帽检测,因此下面我将演示如何在 Kaggle 搜索安全帽检测数据集并在 MMDetection 训练的全流程。
进入 Kaggle 的界面后,我们可以看到左边栏里有一栏 Datasets,这就是 Kaggle 里专门存放数据集的地方。我们点击进去后搜索 Helmet Detection(安全帽检测)就可以看到大量的相关数据集。我们可以找到最多人点赞的数据集进去查阅。
进去后可以看到一些基本的信息,包括标注的类别,文件的大小,整体文件格式和内容等等。从格式上看,这是一个 VOC 格式的标注文件(后缀名为 xml),假如我们遇到标注文件为 txt 的大概就是 YOLO 格式的标注文件,假如我们遇到 json 格式的标注文件那大概就是 COCO 和 JSON 格式的标注文件,不同的标注文件都有不同的格式,我们应用的时候都需要根据具体的要求进行转换。
假如我们认为这部分数据有意义的话,我们就可以下载到 MMDetection 下的 data 文件夹里(没有的话就新建一个)进行保存。
这样我们的数据准备就完成了,下面我们就需要将这部分数据转换成适合于模型训练的数据并开始模型的训练。
格式转换
然后假如我们要了解如何调用该模型以及我们的原数据该准备怎么样的结构输入进去的话,我们可以利用大语言模型的解析能力来为我们进行分析(下面是靠 ChatGPT 解析出来的结果)。
提示词:
#上面是代码的内容
请你作为一名 python 的代码专家,请你根据上述的代码,为我进行分析我所需要准备的数据是什么结构的,然后如何进行调用这个代码呢?
因此下面我们就需要调整一下我们原有的文件格式。下面是原有的数据结构。
可以看到现在的文件格式和所需要的文件格式并不匹配,我们需要进一步调整该数据结构才能使用官方配备的转换代码进行转换。这部分的工作其实很重复枯燥,我们同样可以利用大语言模型来创建脚本辅助我们完成。
提示词:
请你作为一名 python 的脚本专家,辅助我完成以下任务。在 data/archive 文件夹下有两个文件夹。一个是 annotations,里面存放着所有的 xml 标注文件。另一个是 images,里面存放了所有的原始图片文件。现在我希望你能够帮我写一个脚本完成一些重复的操作。
首先你先把 archive 文件夹改名为 VOC2007,并且把里面的 annotations 文件的名称修改为 Annotations,同时把 images 的文件夹名称修改为 JPEGImages。并且检查 JPEGImages 文件夹里的文件是否为 jpg 文件,假如里面的图片文件不是 jpg 文件,请统一将里面的图片文件修改为 jpg 格式。
下面是利用大语言模型 ChatGPT 所获得的的回复。
好的,我们可以分几个步骤来完成这个任务:
重命名文件夹:将 archive 重命名为 VOC2007,并在内部将 annotations 重命名为 Annotations,images 重命名为 JPEGImages。
转换图片格式:检查 JPEGImages 文件夹里的所有文件,如果发现非 .jpg 文件,转换为 .jpg 格式。
创建文件夹结构:在 data/VOC2007 下创建 ImageSets/Main 文件夹结构。
以下是完成上述步骤的 Python 脚本代码:
在运行脚本之前,请确保目标目录下有 data/archive 文件夹存在,且该文件夹下有 annotations 和 images 两个子文件夹。另外,请确保已经安装了 Pillow 库,如果没有安装,请使用 pip install Pillow 进行安装。
您可以将上述代码保存为 .py 文件,并在包含 data/archive 的目录中运行它。完成上述操作后,应会在控制台中看到"已完成"的消息。
我们运行脚本后就能得到以下格式的文件:
接下来我们就可以调用脚本来进行格式的转换了,调用脚本的格式如下(之所以输出格式是这样主要是满足数据集制定的要求):
python:这是启动 Python 解释器的命令。
data: 这是命令行参数 devkit_path,表示 VOC 数据集开发工具包的路径。它应该包含 VOC2007 和/或 VOC2012 数据集的文件夹。
--out-dir data/coco/annotations:这是一个可选参数,用于指定输出目录。在这里,转换后的COCO格式的JSON文件将被保存在 data/coco/annotations 目录下。
--out-format coco:这个可选参数定义了输出格式,这里指定为 COCO 格式。
在 annotations 里创建了 train 和 val 两个 COCO 格式的标注文件,在 train2017 和 val2017 下按照具体数据集需要的格式创建了以下文件夹,并且 JPEGImages 里保存了所有训练集的图片和验证集的图片。这样我们的数据就已经符合 MMDetection 的训练条件了,下面我们可以创建自己的数据集类以及修改配置文件来真正的开始模型的训练了!
创建新的数据集类
首先我们需要先进入 mmdet/datasets 这个文件夹里,并创建一个自己命名数据集的 Python 文件,比如说 helmet。
下面这张图是原本 CocoDataset 的内容:
下面这张图是修改为 HelmetDataset 的内容:
配置文件准备及修改
由于在 MMDetection 或者 MMYOLO 里大多用的都是 COCO 格式的数据集,因此我认为在这里我们直接用终端进行调用然后修改 work_dir 文件夹里对应的配置文件会比起重新创建更加的方便。
首先我们需要找到 config 文件夹里我们想要的模型文件(我常用的就是 yolo),然后查看里面的 README 文件来看适合我们显卡内存的配置文件(我的显存有 8G,所以基本都可以使用)。越大的模型效果越好,但是假如我们训练的时候出现 CUDA out of memory 的话,那就是显存不够或者算力不够的情况,我们就需要选择更小的模型进行训练。
然后选择了一个 mobilenetv2 的模型直接在终端输入以下命令。
然后我们点击配置文件,为了让模型能够顺利训练,我们主要修改的就是 dataset_type 以及 num_classes(数据格式已经在前面转换中完成),我们可以点击 Ctrl+F 找到这部分的内容并进行修改。
把类别从 80 修改为我们的类别 3
将 dataset_type 修改为我们刚刚制作的 HelmetDataset
除此之外,我们其实也可以调整一些模型的超参数。比如说我们数据集很少的话,然后我希望其能够不那么快拟合,我们可以把学习率调得比较小,从而让其能够慢慢的学习到图像的特征。
lr 学习率可以从 0.003 调整为 0.0001
我们还可以修改优化器类型,比如 SGD 修改为 Adam等 。但是一般来说我们其实只有在模型效果不好的时候才会考虑使用这部分操作。
修改完成后,我们就可以调用这个配置文件进行模型的训练了。
模型训练
在完成配置文件后,我们只需要利用终端调用就可以进行训练了。
但是我们输入了代码后,会很惊奇地发现出现了一个问题,就是出现了 need at least one array to concatenate 的报错。这个报错其实比较常见,其主要原因是数据集或者是配置文件出了问题。在检查了自己的数据集并查阅了网上的解决方案后,我定位了问题。就是我们需要以类似于创建 Datasets 的方式创建一个 metainfo,将自己的 classes(标签类别)和 palette(调色板)写入,并且在 train_dataloader、val_dataloader 和 test_dataloader 里写入即可。
创建 metainfo
写入 train_dataloader
写入 val_dataloader
写入 test_dataloader
修改后我们再尝试在终端运行代码。
训练过程中常见的问题及方法:
假如出现了 CUDA out of memory,这个主要的原因就是在于显卡的显存不足。这个时候就需要选择更小的模型或者是替换更高算力的卡进行使用。
假如出现了上面这个 need at least one array to concatenate 的报错,就要按照我上面讲的方式进行处理。这个报错其实在 YOLO 模型中会经常出现,我在调用 YOLOV5 模型的时候就因为这个问题困扰了很久。
模型结果
当训练完成后,训练好的结果就能够在 work_dir 里查看到,里面有模型的配置文件,权重文件以及训练过程中的记录。
--keys loss:指定了要绘制的关键指标,这里是损失(loss)。
--legend loss:为绘制的曲线指定图例。
打印出来的 loss 图如下所示:
可以看到 loss 随着训练逐步降低,从原本的 2000 降到接近于 300。但是这个还是一个相对较高的 loss,因此模型的准确度也仅在 30 %左右。我们可以在网上实际找几张图片来测试一下模型在实际图片上的效果如何。
模型测试
我在网上随意找了几张图片放进去测试。切记不能够在原本数据集找图片进行测试,除非是预留了 test。因为被训练的图像都已经过拟合了,因此其就算判断正确不也能够确保是模型的效果很好还是过拟合的程度很高。
得到的结果为:
这里虽然预测结果显示为 person,但是这并不代表模型检测真的认为是 person,可能只是因为输出的类是 person 而已。这里的 person 应该就是代表的是 helmet,我们后期只需要对其更改一下即可。但是总体来看,模型的效果还是可以的,即便是有多个人站在一起或者说是在一些复杂的施工场地下也能够识别出来带了安全帽的人,因此该模型还是有可取之处的。
总结
MMDetection 不仅提供了丰富的模型选择和配置,还提供了完善的安装、配置和训练流程。用户可以根据自己的需求,从中选择合适的模型进行训练和推理,无论是追求速度还是效果,MMDetection 都能提供相应的解决方案。
此外,MMDetection 还为用户提供了强大的可视化工具和日志分析工具,使得用户可以直观地观察模型的训练过程和结果,及时调整策略。通过充分利用 MMDetection 的这些工具和特性,即使是深度学习领域的新手也能有效地进行模型训练和优化,推进自己的研究或项目。
相信经过以上的详细介绍,大家已经对 MMDetection 工具箱有了深入的认识,那么是时候开始一步步训练一个自己的 MMDetection 模型啦。