里中信息门户网

 当前位置:里中信息门户网»时事»同声国际娱乐网址·如何用 coco 数据集训练 Detectron2 模型?

同声国际娱乐网址·如何用 coco 数据集训练 Detectron2 模型?

2020-01-11 15:34:52 | 里中信息门户网

同声国际娱乐网址·如何用 coco 数据集训练 Detectron2 模型?

同声国际娱乐网址,随着最新的 pythorc1.3 版本的发布,下一代完全重写了它以前的目标检测框架,新的目标检测框架被称为 detectron2。本教程将通过使用自定义 coco 数据集训练实例分割模型,帮助你开始使用此框架。如果你不知道如何创建 coco 数据集,请阅读我之前的文章——如何创建自定义 coco 数据集用于实例分割。

为了快速开始,我们将在 colab notebook 上进行实验,这样你就不必担心在使用 pytorch 1.3 和 detectron2 之前在自己的机器上设置开发环境的问题了。

安装 detectron2

在 colab notebook 中,只需运行这 4 行代码即可安装最新的 pytorch 1.3 和 detectron2。

!pip install -u torch torchvision

!pip install git+https://github.com/facebookresearch/fvcore.git

!git clone https://github.com/facebookresearch/detectron2 detectron2_repo

!pip install -e detectron2_repo

单击输出单元格中的「restart runtime」以使安装生效。

注册一个 coco 数据集

为了告诉 detectron2 如何获取数据集,我们将「注册」它。

为了演示这个过程,我们使用了水果坚果分割数据集,它只有 3 个类:数据、图和榛子。我们将从现有的 coco 数据集训练模型中分离出一个分割模型,该模型可在 decttron2 model zoo 中使用。

你可以这样下载数据集。

# download, decompress the data

!wget https://github.com/tony607/detectron2_instance_segmentation_demo/releases/download/v0.1/data

!unzip data.zip > /dev/null

或者你也可以从这里上传你自己的数据集。

按照 detectron2 自定义数据集教程,将水果坚果数据集注册到 detectron2。

from detectron2.data.datasets import

register_coco_instances register_coco_instances("fruits_nuts", {}, "./data/trainval.json", "./data/images")

每个数据集都与一些元数据相关联。在我们的例子中,可以通过调用fruits_nuts_metadata=metadatacatalog.get(“fruits_nuts”)来访问它。

metadata(evaluator_type='coco', image_root='./data/images',

json_file='./data/trainval.json', name='fruits_nuts', thing_classes=['date', 'fig', 'hazelnut'], thing_dataset_id_to_contiguous_id={1: 0, 2: 1, 3: 2})

要获取目录的实际内部表示形式,可以调用 dataset_dicts=dataset catalog.get("fruits_nuts")。内部格式使用一个 dict 来表示一个图像的注释。

为了验证数据加载是否正确,让我们可视化数据集中随机选择的样本的注释:

import random

from detectron2.utils.visualizer import visualizer

for d in random.sample(dataset_dicts, 3):

img = cv2.imread(d["file_name"])

visualizer = visualizer(img[:, :, ::-1], metadata=fruits_nuts_metadata, scale=0.5)

vis = visualizer.draw_dataset_dict(d)

cv2_imshow(vis.get_image[:, :, ::-1])

其中一张图像可能是这样子的:

模型训练

现在,让我们微调水果坚果数据集上的 coco 预训练 r50-fpn mask r-cnn 模型。在 colab 的 k80 gpu 上训练 300 次迭代需要大约 6 分钟。

from detectron2.engine import defaulttrainer

from detectron2.config import get_cfg

import os

cfg = get_cfg

cfg.merge_from_file(

"./detectron2_repo/configs/coco-instancesegmentation/mask_rcnn_r_50_fpn_3 x.yaml" )

cfg.datasets.train = ("fruits_nuts",)

cfg.datasets.test = # no metrics implemented for this dataset cfg.dataloader.num_workers = 2

cfg.model.weights = "detectron2://coco-instancesegmentation/mask_rcnn_r_50_fpn_3x/137849600/mode

cfg.solver.ims_per_batch = 2

cfg.solver.base_lr = 0.02

cfg.solver.max_iter = ( 300 ) # 300 iterations seems good enough, but you can certainly train longer

cfg.model.roi_heads.batch_size_per_image = ( 128 ) # faster, and good enough for this toy dataset

cfg.model.roi_heads.num_classes = 3 # 3 classes (data, fig, hazelnut)

os.makedirs(cfg.output_dir, exist_ok=true)

trainer = defaulttrainer(cfg)

trainer.resume_or_load(resume=false)

trainer.train

如果切换到自己的数据集,请相应地更改类数、学习速率或最大迭代次数。

作出预测

现在,我们用训练好的模型在水果坚果数据集上进行推理。首先,让我们使用我们刚刚训练的模型创建一个预测:

cfg.model.weights = os.path.join(cfg.output_dir, "model_final.pth")

cfg.model.roi_heads.score_thresh_test = 0.5 # set the testing threshold for this model

cfg.datasets.test = ("fruits_nuts", )

predictor = defaultpredictor(cfg)

然后,随机选取多个样本对预测结果进行可视化处理。

from detectron2.utils.visualizer import colormode

for d in random.sample(dataset_dicts, 3):

im = cv2.imread(d["file_name"])

outputs = predictor(im)

v = visualizer(im[:, :, ::-1], metadata=fruits_nuts_metadata, scale=0.8, instance_mode=colormode.image_bw # remove the colors of unsegmented pixels )

v = v.draw_instance_predictions(outputs["instances"].to("cpu"))

cv2_imshow(v.get_image[:, :, ::-1])

这是一个覆盖了预测的样本图像所得到的结果。

总结与思考

你可能已经阅读了我以前的教程,其中介绍了一个类似的对象检测框架,名为 mmdetection,它也构建在 pytorch 上。那么 detectron2 和它相比如何呢?以下是我的一些想法。

两个框架都很容易用一个描述模型训练方法的配置文件进行配置。detectron2 的 yaml 配置文件效率更高,有两个原因。首先,可以通过先进行「基本」配置来重用配置,并在此基本配置文件上构建最终的训练配置文件,从而减少重复代码。第二,可以先配置配置文件,并允许在 python 代码中进行必要的进一步修改,从而使其更加灵活。

那么推理速度如何?简单地说,detectron2 比相同 mask rcnn resnet50 fpn 模型的 mmdetection 稍快。mmdetection 的 fps 是 2.45,而 detectron2 达到 2.59 fps,在推断单个图像时提高了 5.7% 的速度。我们基于以下代码做了基准测试。

import time

times =

for i in range(20):

start_time = time.time

outputs = predictor(im)

delta = time.time - start_time

times.append(delta)

mean_delta = np.array(times).mean

fps = 1 / mean_delta

print("average(sec):{:.2f},fps:{:.2f}".format(mean_delta, fps))

所以,你现在学会啦,detectron2 让你用自定义数据集训练自定义实例分割模型变得非常简单。你可能会发现以下资源很有帮助:

我之前的文章——how to create custom coco data set for instance segmentation。

我之前的文章——how to train an object detection model with mmdetection。

detectron2 github repo。

这篇文章的可运行的 colab notebook。

via:https://medium.com/@chengweizhang2012/how-to-train-detectron2-with-custom-coco-datasets-4d5170c9f389

雷锋网雷锋网雷锋网

新闻

栏目资讯

推荐

Copyright 2018-2019 avnile.com 里中信息门户网 Inc. All Rights Reserved.