Detectron2快速入门(2)

滴滴云技术支持发表于:2020年11月09日 11:56:06更新于:2020年11月09日 12:00:28

在本教程中,您将学习如何使用Detectron2 在自定义数据集上训练一个目标检测网络。我们需要下载一个非coco格式的目标检测数据集

!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
# 我们提前下载好,也可通过滴滴云S3下载
# !wget https://dataset-public.s3-internal.didiyunapi.com/DAI教程/datasets/balloon_dataset.zip
!unzip balloon_dataset.zip
--2020-04-26 15:14:09--  https://dataset-public.s3-internal.didiyunapi.com/DAI%E6%95%99%E7%A8%8B/datasets/balloon_dataset.zip
Resolving dataset-public.s3-internal.didiyunapi.com (dataset-public.s3-internal.didiyunapi.com)... 100.64.8.5
Connecting to dataset-public.s3-internal.didiyunapi.com (dataset-public.s3-internal.didiyunapi.com)|100.64.8.5|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 38741381 (37M) [application/zip]
Saving to: ‘balloon_dataset.zip’

100%[======================================>] 38,741,381   110MB/s   in 0.3s   

2020-04-26 15:14:09 (110 MB/s) - ‘balloon_dataset.zip’ saved [38741381/38741381]
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

import numpy as np
import cv2
import random

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog

import matplotlib.pyplot as plt
def cv2_imshow(img_bgr):
    plt.rcParams['figure.figsize'] = (18, 36)
    plt.axis('off')
    plt.imshow(img_bgr[...,::-1])
# 下面我们需要把自定义数据格式改成标准格式,如果数据集为标准的COCO格式,可以省去格式处理代码,用下面3行代替即可
# from detectron2.data.datasets import register_coco_instances
# register_coco_instances("my_dataset_train", {}, "json_annotation_train.json", "path/to/image/dir")
# register_coco_instances("my_dataset_val", {}, "json_annotation_val.json", "path/to/image/dir")

import os
import json
from detectron2.structures import BoxMode

def get_balloon_dicts(img_dir):
    json_file = os.path.join(img_dir, "via_region_data.json")
    with open(json_file) as f:
        imgs_anns = json.load(f)

    dataset_dicts = []
    for idx, v in enumerate(imgs_anns.values()):
        record = {}
        
        filename = os.path.join(img_dir, v["filename"])
        height, width = cv2.imread(filename).shape[:2]
        
        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width
      
        annos = v["regions"]
        objs = []
        for _, anno in annos.items():
            assert not anno["region_attributes"]
            anno = anno["shape_attributes"]
            px = anno["all_points_x"]
            py = anno["all_points_y"]
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]

            obj = {
                "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
                "bbox_mode": BoxMode.XYXY_ABS,
                "segmentation": [poly],
                "category_id": 0,
                "iscrowd": 0
            }
            objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts

from detectron2.data import DatasetCatalog, MetadataCatalog
for d in ["train", "val"]:
    DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
    MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")
# 我们从数据集里随机可视化一张图片和标注,来验证数据加载是否正确
    
dataset_dicts = get_balloon_dicts("balloon/train")
for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    cv2_imshow(vis.get_image()[:, :, ::-1])

0015fa8bea996c667109058490cb80f

# 下面我们使用预训练模型 R50-FPN Mask R-CNN 开始训练

from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = cfg.MODEL.WEIGHTS = '/home/dc2-user/Detectron2/model_final_b275ba.pkl' # 请参考(1)下载预训练模型
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.0001
cfg.SOLVER.MAX_ITER = 500
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()
•[32m[04/26 15:35:02 d2.engine.defaults]: •[0mModel:
GeneralizedRCNN(
  (backbone): FPN(
    (fpn_lateral2): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral3): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral4): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral5): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (top_block): LastLevelMaxPool()
    (bottom_up): ResNet(
      (stem): BasicStem(
        (conv1): Conv2d(
          3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
          (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
        )
      )
      (res2): Sequential(
        (0): BottleneckBlock(
          (shortcut): Conv2d(
            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv1): Conv2d(
            64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv2): Conv2d(
            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv3): Conv2d(
            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
        )
        (1): BottleneckBlock(
          (conv1): Conv2d(
            256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv2): Conv2d(
            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv3): Conv2d(
            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
        )
        (2): BottleneckBlock(
          (conv1): Conv2d(
            256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv2): Conv2d(
            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)
          )
          (conv3): Conv2d(
            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
        )
      )
      (res3): Sequential(
        (0): BottleneckBlock(
          (shortcut): Conv2d(
            256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv1): Conv2d(
            256, 128, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv2): Conv2d(
            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv3): Conv2d(
            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
        )
        (1): BottleneckBlock(
          (conv1): Conv2d(
            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv2): Conv2d(
            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv3): Conv2d(
            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
        )
        (2): BottleneckBlock(
          (conv1): Conv2d(
            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv2): Conv2d(
            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv3): Conv2d(
            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
        )
        (3): BottleneckBlock(
          (conv1): Conv2d(
            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv2): Conv2d(
            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)
          )
          (conv3): Conv2d(
            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
        )
      )
      (res4): Sequential(
        (0): BottleneckBlock(
          (shortcut): Conv2d(
            512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
          (conv1): Conv2d(
            512, 256, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
        (1): BottleneckBlock(
          (conv1): Conv2d(
            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
        (2): BottleneckBlock(
          (conv1): Conv2d(
            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
        (3): BottleneckBlock(
          (conv1): Conv2d(
            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
        (4): BottleneckBlock(
          (conv1): Conv2d(
            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
        (5): BottleneckBlock(
          (conv1): Conv2d(
            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv2): Conv2d(
            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)
          )
          (conv3): Conv2d(
            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)
          )
        )
      )
      (res5): Sequential(
        (0): BottleneckBlock(
          (shortcut): Conv2d(
            1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)
          )
          (conv1): Conv2d(
            1024, 512, kernel_size=(1, 1), stride=(2, 2), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv2): Conv2d(
            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv3): Conv2d(
            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)
          )
        )
        (1): BottleneckBlock(
          (conv1): Conv2d(
            2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv2): Conv2d(
            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv3): Conv2d(
            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)
          )
        )
        (2): BottleneckBlock(
          (conv1): Conv2d(
            2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv2): Conv2d(
            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)
          )
          (conv3): Conv2d(
            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)
          )
        )
      )
    )
  )
  (proposal_generator): RPN(
    (anchor_generator): DefaultAnchorGenerator(
      (cell_anchors): BufferList()
    )
    (rpn_head): StandardRPNHead(
      (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (objectness_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
      (anchor_deltas): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
    )
  )
  (roi_heads): StandardROIHeads(
    (box_pooler): ROIPooler(
      (level_poolers): ModuleList(
        (0): ROIAlign(output_size=(7, 7), spatial_scale=0.25, sampling_ratio=0, aligned=True)
        (1): ROIAlign(output_size=(7, 7), spatial_scale=0.125, sampling_ratio=0, aligned=True)
        (2): ROIAlign(output_size=(7, 7), spatial_scale=0.0625, sampling_ratio=0, aligned=True)
        (3): ROIAlign(output_size=(7, 7), spatial_scale=0.03125, sampling_ratio=0, aligned=True)
      )
    )
    (box_head): FastRCNNConvFCHead(
      (fc1): Linear(in_features=12544, out_features=1024, bias=True)
      (fc2): Linear(in_features=1024, out_features=1024, bias=True)
    )
    (box_predictor): FastRCNNOutputLayers(
      (cls_score): Linear(in_features=1024, out_features=2, bias=True)
      (bbox_pred): Linear(in_features=1024, out_features=4, bias=True)
    )
    (mask_pooler): ROIPooler(
      (level_poolers): ModuleList(
        (0): ROIAlign(output_size=(14, 14), spatial_scale=0.25, sampling_ratio=0, aligned=True)
        (1): ROIAlign(output_size=(14, 14), spatial_scale=0.125, sampling_ratio=0, aligned=True)
        (2): ROIAlign(output_size=(14, 14), spatial_scale=0.0625, sampling_ratio=0, aligned=True)
        (3): ROIAlign(output_size=(14, 14), spatial_scale=0.03125, sampling_ratio=0, aligned=True)
      )
    )
    (mask_head): MaskRCNNConvUpsampleHead(
      (mask_fcn1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (mask_fcn2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (mask_fcn3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (mask_fcn4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (deconv): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))
      (predictor): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))
    )
  )
)
•[32m[04/26 15:35:04 d2.data.build]: •[0mRemoved 0 images with no usable annotations. 61 images left.
•[32m[04/26 15:35:04 d2.data.common]: •[0mSerializing 61 elements to byte tensors and concatenating them all ...
•[32m[04/26 15:35:04 d2.data.common]: •[0mSerialized dataset takes 0.17 MiB
•[32m[04/26 15:35:04 d2.data.detection_utils]: •[0mTransformGens used in training: [ResizeShortestEdge(short_edge_length=(640, 672, 704, 736, 768, 800), max_size=1333, sample_style='choice'), RandomFlip()]
•[32m[04/26 15:35:04 d2.data.build]: •[0mUsing training sampler TrainingSampler
'roi_heads.box_predictor.cls_score.weight' has shape (81, 1024) in the checkpoint but (2, 1024) in the model! Skipped.
'roi_heads.box_predictor.cls_score.bias' has shape (81,) in the checkpoint but (2,) in the model! Skipped.
'roi_heads.box_predictor.bbox_pred.weight' has shape (320, 1024) in the checkpoint but (4, 1024) in the model! Skipped.
'roi_heads.box_predictor.bbox_pred.bias' has shape (320,) in the checkpoint but (4,) in the model! Skipped.
•[32m[04/26 15:35:05 d2.engine.train_loop]: •[0mStarting training from iteration 0
•[32m[04/26 15:35:19 d2.utils.events]: •[0m eta: 0:04:26  iter: 19  total_loss: 2.130  loss_cls: 0.735  loss_box_reg: 0.640  loss_mask: 0.690  loss_rpn_cls: 0.031  loss_rpn_loc: 0.007  time: 0.5656  data_time: 0.0406  lr: 0.000002  max_mem: 2724M
•[32m[04/26 15:35:31 d2.utils.events]: •[0m eta: 0:04:16  iter: 39  total_loss: 2.083  loss_cls: 0.725  loss_box_reg: 0.633  loss_mask: 0.690  loss_rpn_cls: 0.031  loss_rpn_loc: 0.009  time: 0.5661  data_time: 0.0088  lr: 0.000004  max_mem: 2724M
•[32m[04/26 15:35:42 d2.utils.events]: •[0m eta: 0:04:06  iter: 59  total_loss: 2.107  loss_cls: 0.694  loss_box_reg: 0.615  loss_mask: 0.689  loss_rpn_cls: 0.041  loss_rpn_loc: 0.011  time: 0.5671  data_time: 0.0093  lr: 0.000006  max_mem: 2724M
•[32m[04/26 15:35:54 d2.utils.events]: •[0m eta: 0:03:58  iter: 79  total_loss: 2.026  loss_cls: 0.656  loss_box_reg: 0.645  loss_mask: 0.686  loss_rpn_cls: 0.032  loss_rpn_loc: 0.011  time: 0.5713  data_time: 0.0089  lr: 0.000008  max_mem: 2724M
•[32m[04/26 15:36:05 d2.utils.events]: •[0m eta: 0:03:47  iter: 99  total_loss: 1.989  loss_cls: 0.601  loss_box_reg: 0.637  loss_mask: 0.685  loss_rpn_cls: 0.043  loss_rpn_loc: 0.010  time: 0.5746  data_time: 0.0085  lr: 0.000010  max_mem: 2846M
•[32m[04/26 15:36:17 d2.utils.events]: •[0m eta: 0:03:37  iter: 119  total_loss: 1.952  loss_cls: 0.551  loss_box_reg: 0.595  loss_mask: 0.682  loss_rpn_cls: 0.033  loss_rpn_loc: 0.010  time: 0.5761  data_time: 0.0086  lr: 0.000012  max_mem: 2846M
•[32m[04/26 15:36:29 d2.utils.events]: •[0m eta: 0:03:29  iter: 139  total_loss: 1.828  loss_cls: 0.518  loss_box_reg: 0.586  loss_mask: 0.678  loss_rpn_cls: 0.021  loss_rpn_loc: 0.005  time: 0.5791  data_time: 0.0085  lr: 0.000014  max_mem: 2846M
•[32m[04/26 15:36:41 d2.utils.events]: •[0m eta: 0:03:18  iter: 159  total_loss: 1.868  loss_cls: 0.484  loss_box_reg: 0.709  loss_mask: 0.675  loss_rpn_cls: 0.034  loss_rpn_loc: 0.010  time: 0.5804  data_time: 0.0087  lr: 0.000016  max_mem: 2846M
•[32m[04/26 15:36:53 d2.utils.events]: •[0m eta: 0:03:06  iter: 179  total_loss: 1.836  loss_cls: 0.443  loss_box_reg: 0.653  loss_mask: 0.672  loss_rpn_cls: 0.034  loss_rpn_loc: 0.008  time: 0.5808  data_time: 0.0091  lr: 0.000018  max_mem: 2846M
•[32m[04/26 15:37:04 d2.utils.events]: •[0m eta: 0:02:55  iter: 199  total_loss: 1.824  loss_cls: 0.423  loss_box_reg: 0.694  loss_mask: 0.666  loss_rpn_cls: 0.020  loss_rpn_loc: 0.006  time: 0.5814  data_time: 0.0090  lr: 0.000020  max_mem: 2846M
•[32m[04/26 15:37:16 d2.utils.events]: •[0m eta: 0:02:43  iter: 219  total_loss: 1.597  loss_cls: 0.379  loss_box_reg: 0.504  loss_mask: 0.663  loss_rpn_cls: 0.038  loss_rpn_loc: 0.007  time: 0.5816  data_time: 0.0092  lr: 0.000022  max_mem: 2846M
•[32m[04/26 15:37:28 d2.utils.events]: •[0m eta: 0:02:32  iter: 239  total_loss: 1.698  loss_cls: 0.362  loss_box_reg: 0.645  loss_mask: 0.656  loss_rpn_cls: 0.028  loss_rpn_loc: 0.008  time: 0.5815  data_time: 0.0082  lr: 0.000024  max_mem: 2846M
•[32m[04/26 15:37:40 d2.utils.events]: •[0m eta: 0:02:20  iter: 259  total_loss: 1.756  loss_cls: 0.357  loss_box_reg: 0.687  loss_mask: 0.655  loss_rpn_cls: 0.021  loss_rpn_loc: 0.011  time: 0.5831  data_time: 0.0083  lr: 0.000026  max_mem: 2846M
•[32m[04/26 15:37:52 d2.utils.events]: •[0m eta: 0:02:09  iter: 279  total_loss: 1.640  loss_cls: 0.326  loss_box_reg: 0.639  loss_mask: 0.653  loss_rpn_cls: 0.034  loss_rpn_loc: 0.009  time: 0.5843  data_time: 0.0084  lr: 0.000028  max_mem: 2846M
•[32m[04/26 15:38:04 d2.utils.events]: •[0m eta: 0:01:57  iter: 299  total_loss: 1.611  loss_cls: 0.305  loss_box_reg: 0.649  loss_mask: 0.640  loss_rpn_cls: 0.023  loss_rpn_loc: 0.007  time: 0.5852  data_time: 0.0091  lr: 0.000030  max_mem: 2846M
•[32m[04/26 15:38:15 d2.utils.events]: •[0m eta: 0:01:46  iter: 319  total_loss: 1.533  loss_cls: 0.285  loss_box_reg: 0.576  loss_mask: 0.624  loss_rpn_cls: 0.026  loss_rpn_loc: 0.007  time: 0.5846  data_time: 0.0085  lr: 0.000032  max_mem: 2846M
•[32m[04/26 15:38:27 d2.utils.events]: •[0m eta: 0:01:34  iter: 339  total_loss: 1.631  loss_cls: 0.287  loss_box_reg: 0.661  loss_mask: 0.643  loss_rpn_cls: 0.023  loss_rpn_loc: 0.011  time: 0.5858  data_time: 0.0082  lr: 0.000034  max_mem: 2846M
•[32m[04/26 15:38:39 d2.utils.events]: •[0m eta: 0:01:22  iter: 359  total_loss: 1.524  loss_cls: 0.235  loss_box_reg: 0.603  loss_mask: 0.622  loss_rpn_cls: 0.020  loss_rpn_loc: 0.007  time: 0.5854  data_time: 0.0085  lr: 0.000036  max_mem: 2846M
•[32m[04/26 15:38:51 d2.utils.events]: •[0m eta: 0:01:11  iter: 379  total_loss: 1.538  loss_cls: 0.235  loss_box_reg: 0.659  loss_mask: 0.623  loss_rpn_cls: 0.020  loss_rpn_loc: 0.010  time: 0.5856  data_time: 0.0085  lr: 0.000038  max_mem: 2846M
•[32m[04/26 15:39:03 d2.utils.events]: •[0m eta: 0:00:59  iter: 399  total_loss: 1.414  loss_cls: 0.202  loss_box_reg: 0.573  loss_mask: 0.609  loss_rpn_cls: 0.016  loss_rpn_loc: 0.004  time: 0.5861  data_time: 0.0087  lr: 0.000040  max_mem: 2846M
•[32m[04/26 15:39:15 d2.utils.events]: •[0m eta: 0:00:47  iter: 419  total_loss: 1.478  loss_cls: 0.216  loss_box_reg: 0.579  loss_mask: 0.611  loss_rpn_cls: 0.026  loss_rpn_loc: 0.011  time: 0.5863  data_time: 0.0091  lr: 0.000042  max_mem: 2935M
•[32m[04/26 15:39:27 d2.utils.events]: •[0m eta: 0:00:35  iter: 439  total_loss: 1.341  loss_cls: 0.180  loss_box_reg: 0.514  loss_mask: 0.595  loss_rpn_cls: 0.024  loss_rpn_loc: 0.006  time: 0.5870  data_time: 0.0085  lr: 0.000044  max_mem: 2935M
•[32m[04/26 15:39:39 d2.utils.events]: •[0m eta: 0:00:24  iter: 459  total_loss: 1.420  loss_cls: 0.186  loss_box_reg: 0.582  loss_mask: 0.599  loss_rpn_cls: 0.024  loss_rpn_loc: 0.008  time: 0.5873  data_time: 0.0084  lr: 0.000046  max_mem: 2935M
•[32m[04/26 15:39:51 d2.utils.events]: •[0m eta: 0:00:12  iter: 479  total_loss: 1.345  loss_cls: 0.169  loss_box_reg: 0.522  loss_mask: 0.601  loss_rpn_cls: 0.020  loss_rpn_loc: 0.007  time: 0.5877  data_time: 0.0084  lr: 0.000048  max_mem: 2935M
•[32m[04/26 15:40:03 d2.utils.events]: •[0m eta: 0:00:00  iter: 499  total_loss: 1.215  loss_cls: 0.143  loss_box_reg: 0.487  loss_mask: 0.564  loss_rpn_cls: 0.016  loss_rpn_loc: 0.006  time: 0.5887  data_time: 0.0087  lr: 0.000050  max_mem: 2935M
•[32m[04/26 15:40:05 d2.engine.hooks]: •[0mOverall training speed: 497 iterations in 0:04:53 (0.5900 s / it)
•[32m[04/26 15:40:05 d2.engine.hooks]: •[0mTotal training time: 0:04:55 (0:00:02 on hooks)
# 训练完成后我们来测试一下效果
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
cfg.DATASETS.TEST = ("balloon_val", )
predictor = DefaultPredictor(cfg)

from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_balloon_dicts("balloon/val")

for d in random.sample(dataset_dicts, 3):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=balloon_metadata, 
                   scale=0.8, 
                   instance_mode=ColorMode.IMAGE_BW
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(v.get_image()[:, :, ::-1])

png0015fa8becf0c3982a66b9dcb20ae9b



png