隨著最新的 Pythorc1.3 版本的發(fā)布，下一代完全重寫了它以前的目標(biāo)檢測(cè)框架，新的目標(biāo)檢測(cè)框架被稱為 Detectron2。本教程將通過使用自定義 coco 數(shù)據(jù)集訓(xùn)練實(shí)例分割模型，幫助你開始使用此框架。如果你不知道如何創(chuàng)建 coco 數(shù)據(jù)集，請(qǐng)閱讀我之前的文章——如何創(chuàng)建自定義 coco 數(shù)據(jù)集用于實(shí)例分割。

為了快速開始，我們將在 Colab Notebook 上進(jìn)行實(shí)驗(yàn)，這樣你就不必?fù)?dān)心在使用 pytorch 1.3 和 detectron2 之前在自己的機(jī)器上設(shè)置開發(fā)環(huán)境的問題了。

安裝 Detectron2

在 Colab Notebook 中，只需運(yùn)行這 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」以使安裝生效。

注冊(cè)一個(gè) coco 數(shù)據(jù)集

為了告訴 Detectron2 如何獲取數(shù)據(jù)集，我們將「注冊(cè)」它。

為了演示這個(gè)過程，我們使用了水果堅(jiān)果分割數(shù)據(jù)集，它只有 3 個(gè)類：數(shù)據(jù)、圖和榛子。我們將從現(xiàn)有的 coco 數(shù)據(jù)集訓(xùn)練模型中分離出一個(gè)分割模型，該模型可在 DeCtTrON2 model zoo 中使用。

你可以這樣下載數(shù)據(jù)集。

# download, decompress the data

!wget https://github.com/Tony607/detectron2_instance_segmentation_demo/releases/download/V0.1/data

!unzip data.zip > /dev/null

或者你也可以從這里上傳你自己的數(shù)據(jù)集。

按照 Detectron2 自定義數(shù)據(jù)集教程，將水果堅(jiān)果數(shù)據(jù)集注冊(cè)到 Detectron2。

from detectron2.data.datasets import

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

每個(gè)數(shù)據(jù)集都與一些元數(shù)據(jù)相關(guān)聯(lián)。在我們的例子中，可以通過調(diào)用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})

要獲取目錄的實(shí)際內(nèi)部表示形式，可以調(diào)用 dataset_dicts=dataset catalog.get（"fruits_nuts"）。內(nèi)部格式使用一個(gè) dict 來表示一個(gè)圖像的注釋。

為了驗(yàn)證數(shù)據(jù)加載是否正確，讓我們可視化數(shù)據(jù)集中隨機(jī)選擇的樣本的注釋：

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])

其中一張圖像可能是這樣子的：

模型訓(xùn)練

現(xiàn)在，讓我們微調(diào)水果堅(jiān)果數(shù)據(jù)集上的 coco 預(yù)訓(xùn)練 R50-FPN Mask R-CNN 模型。在 colab 的 k80 gpu 上訓(xùn)練 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()

如果切換到自己的數(shù)據(jù)集，請(qǐng)相應(yīng)地更改類數(shù)、學(xué)習(xí)速率或最大迭代次數(shù)。

作出預(yù)測(cè)

現(xiàn)在，我們用訓(xùn)練好的模型在水果堅(jiān)果數(shù)據(jù)集上進(jìn)行推理。首先，讓我們使用我們剛剛訓(xùn)練的模型創(chuàng)建一個(gè)預(yù)測(cè)：

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)

然后，隨機(jī)選取多個(gè)樣本對(duì)預(yù)測(cè)結(jié)果進(jìn)行可視化處理。

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])

這是一個(gè)覆蓋了預(yù)測(cè)的樣本圖像所得到的結(jié)果。

總結(jié)與思考

你可能已經(jīng)閱讀了我以前的教程，其中介紹了一個(gè)類似的對(duì)象檢測(cè)框架，名為 MMdetection，它也構(gòu)建在 pytorch 上。那么 Detectron2 和它相比如何呢？以下是我的一些想法。

兩個(gè)框架都很容易用一個(gè)描述模型訓(xùn)練方法的配置文件進(jìn)行配置。Detectron2 的 yaml 配置文件效率更高，有兩個(gè)原因。首先，可以通過先進(jìn)行「基本」配置來重用配置，并在此基本配置文件上構(gòu)建最終的訓(xùn)練配置文件，從而減少重復(fù)代碼。第二，可以先配置配置文件，并允許在 Python 代碼中進(jìn)行必要的進(jìn)一步修改，從而使其更加靈活。

那么推理速度如何？簡(jiǎn)單地說，Detectron2 比相同 Mask RCNN Resnet50 FPN 模型的 MMdetection 稍快。MMdetection 的 FPS 是 2.45，而 Detectron2 達(dá)到 2.59 FPS，在推斷單個(gè)圖像時(shí)提高了 5.7% 的速度。我們基于以下代碼做了基準(zhǔn)測(cè)試。

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))

所以，你現(xiàn)在學(xué)會(huì)啦，Detectron2 讓你用自定義數(shù)據(jù)集訓(xùn)練自定義實(shí)例分割模型變得非常簡(jiǎn)單。你可能會(huì)發(fā)現(xiàn)以下資源很有幫助：

• 我之前的文章——How to create custom COCO data set for instance segmentation。
  

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

• Detectron2 GitHub repo。
  

• 這篇文章的可運(yùn)行的 Colab Notebook 。
  

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

雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)

雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。