ショートカット

モデルのカスタマイズ

基本的に、モデルのコンポーネントを5つのタイプに分類します。

  • バックボーン:通常、特徴マップを抽出するFCNネットワーク(例:ResNet、MobileNet)。

  • ネック:バックボーンとヘッドの間のコンポーネント(例:FPN、PAFPN)。

  • ヘッド:特定のタスク(例:bbox予測とマスク予測)のためのコンポーネント。

  • RoIエクストラクター:特徴マップからRoI特徴を抽出する部分(例:RoI Align)。

  • 損失:損失を計算するためのヘッド内のコンポーネント(例:FocalLoss、L1Loss、GHMLoss)。

新しいコンポーネントの開発

新しいバックボーンの追加

ここでは、MobileNetの例を用いて新しいコンポーネントを開発する方法を示します。

1. 新しいバックボーンを定義する(例:MobileNet)

新しいファイル mmdet/models/backbones/mobilenet.py を作成します。

import torch.nn as nn

from mmdet.registry import MODELS


@MODELS.register_module()
class MobileNet(nn.Module):

    def __init__(self, arg1, arg2):
        pass

    def forward(self, x):  # should return a tuple
        pass

2. モジュールのインポート

次の行を mmdet/models/backbones/__init__.py に追加するか、

from .mobilenet import MobileNet

または、元のコードを変更しないように、

custom_imports = dict(
    imports=['mmdet.models.backbones.mobilenet'],
    allow_failed_imports=False)

を設定ファイルに追加することもできます。

3. 設定ファイルでバックボーンを使用する

model = dict(
    ...
    backbone=dict(
        type='MobileNet',
        arg1=xxx,
        arg2=xxx),
    ...

新しいネックの追加

1. ネックを定義する(例:PAFPN)

新しいファイル mmdet/models/necks/pafpn.py を作成します。

import torch.nn as nn

from mmdet.registry import MODELS

@MODELS.register_module()
class PAFPN(nn.Module):

    def __init__(self,
                in_channels,
                out_channels,
                num_outs,
                start_level=0,
                end_level=-1,
                add_extra_convs=False):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

2. モジュールのインポート

次の行を mmdet/models/necks/__init__.py に追加するか、

from .pafpn import PAFPN

または、元のコードを変更しないように、

custom_imports = dict(
    imports=['mmdet.models.necks.pafpn'],
    allow_failed_imports=False)

を設定ファイルに追加して、元のコードを変更しないようにします。

3. 設定ファイルの変更

neck=dict(
    type='PAFPN',
    in_channels=[256, 512, 1024, 2048],
    out_channels=256,
    num_outs=5)

新しいヘッドの追加

ここでは、以下の例として、Double Head R-CNNの例を用いて、新しいヘッドを開発する方法を示します。

まず、mmdet/models/roi_heads/bbox_heads/double_bbox_head.py に新しいbboxヘッドを追加します。Double Head R-CNNは、物体検出用の新しいbboxヘッドを実装します。bboxヘッドを実装するには、基本的に、新しいモジュールの3つの関数を次のように実装する必要があります。

from typing import Tuple

import torch.nn as nn
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule, ModuleList
from torch import Tensor

from mmdet.models.backbones.resnet import Bottleneck
from mmdet.registry import MODELS
from mmdet.utils import ConfigType, MultiConfig, OptConfigType, OptMultiConfig
from .bbox_head import BBoxHead

@MODELS.register_module()
class DoubleConvFCBBoxHead(BBoxHead):
    r"""Bbox head used in Double-Head R-CNN

    .. code-block:: none

                                          /-> cls
                      /-> shared convs ->
                                          \-> reg
        roi features
                                          /-> cls
                      \-> shared fc    ->
                                          \-> reg
    """  # noqa: W605

    def __init__(self,
                 num_convs: int = 0,
                 num_fcs: int = 0,
                 conv_out_channels: int = 1024,
                 fc_out_channels: int = 1024,
                 conv_cfg: OptConfigType = None,
                 norm_cfg: ConfigType = dict(type='BN'),
                 init_cfg: MultiConfig = dict(
                     type='Normal',
                     override=[
                         dict(type='Normal', name='fc_cls', std=0.01),
                         dict(type='Normal', name='fc_reg', std=0.001),
                         dict(
                             type='Xavier',
                             name='fc_branch',
                             distribution='uniform')
                     ]),
                 **kwargs) -> None:
        kwargs.setdefault('with_avg_pool', True)
        super().__init__(init_cfg=init_cfg, **kwargs)

    def forward(self, x_cls: Tensor, x_reg: Tensor) -> Tuple[Tensor]:

次に、必要に応じて新しいRoI Headを実装します。新しい DoubleHeadRoIHeadStandardRoIHead から継承する予定です。 StandardRoIHead はすでに次の関数を実装していることがわかります。

from typing import List, Optional, Tuple

import torch
from torch import Tensor

from mmdet.registry import MODELS, TASK_UTILS
from mmdet.structures import DetDataSample
from mmdet.structures.bbox import bbox2roi
from mmdet.utils import ConfigType, InstanceList
from ..task_modules.samplers import SamplingResult
from ..utils import empty_instances, unpack_gt_instances
from .base_roi_head import BaseRoIHead


@MODELS.register_module()
class StandardRoIHead(BaseRoIHead):
    """Simplest base roi head including one bbox head and one mask head."""

    def init_assigner_sampler(self) -> None:

    def init_bbox_head(self, bbox_roi_extractor: ConfigType,
                       bbox_head: ConfigType) -> None:

    def init_mask_head(self, mask_roi_extractor: ConfigType,
                       mask_head: ConfigType) -> None:

    def forward(self, x: Tuple[Tensor],
                rpn_results_list: InstanceList) -> tuple:

    def loss(self, x: Tuple[Tensor], rpn_results_list: InstanceList,
             batch_data_samples: List[DetDataSample]) -> dict:

    def _bbox_forward(self, x: Tuple[Tensor], rois: Tensor) -> dict:

    def bbox_loss(self, x: Tuple[Tensor],
                  sampling_results: List[SamplingResult]) -> dict:

    def mask_loss(self, x: Tuple[Tensor],
                  sampling_results: List[SamplingResult], bbox_feats: Tensor,
                  batch_gt_instances: InstanceList) -> dict:

    def _mask_forward(self,
                      x: Tuple[Tensor],
                      rois: Tensor = None,
                      pos_inds: Optional[Tensor] = None,
                      bbox_feats: Optional[Tensor] = None) -> dict:

    def predict_bbox(self,
                     x: Tuple[Tensor],
                     batch_img_metas: List[dict],
                     rpn_results_list: InstanceList,
                     rcnn_test_cfg: ConfigType,
                     rescale: bool = False) -> InstanceList:

    def predict_mask(self,
                     x: Tuple[Tensor],
                     batch_img_metas: List[dict],
                     results_list: InstanceList,
                     rescale: bool = False) -> InstanceList:

Double Headの変更点は主に bbox_forward ロジックであり、他のロジックは StandardRoIHead から継承します。mmdet/models/roi_heads/double_roi_head.py では、次のように新しいRoI Headを実装します。

from typing import Tuple

from torch import Tensor

from mmdet.registry import MODELS
from .standard_roi_head import StandardRoIHead


@MODELS.register_module()
class DoubleHeadRoIHead(StandardRoIHead):
    """RoI head for `Double Head RCNN <https://arxiv.org/abs/1904.06493>`_.

    Args:
        reg_roi_scale_factor (float): The scale factor to extend the rois
            used to extract the regression features.
    """

    def __init__(self, reg_roi_scale_factor: float, **kwargs):
        super().__init__(**kwargs)
        self.reg_roi_scale_factor = reg_roi_scale_factor

    def _bbox_forward(self, x: Tuple[Tensor], rois: Tensor) -> dict:
        """Box head forward function used in both training and testing.

        Args:
            x (tuple[Tensor]): List of multi-level img features.
            rois (Tensor): RoIs with the shape (n, 5) where the first
                column indicates batch id of each RoI.

        Returns:
             dict[str, Tensor]: Usually returns a dictionary with keys:

                - `cls_score` (Tensor): Classification scores.
                - `bbox_pred` (Tensor): Box energies / deltas.
                - `bbox_feats` (Tensor): Extract bbox RoI features.
        """
        bbox_cls_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs], rois)
        bbox_reg_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs],
            rois,
            roi_scale_factor=self.reg_roi_scale_factor)
        if self.with_shared_head:
            bbox_cls_feats = self.shared_head(bbox_cls_feats)
            bbox_reg_feats = self.shared_head(bbox_reg_feats)
        cls_score, bbox_pred = self.bbox_head(bbox_cls_feats, bbox_reg_feats)

        bbox_results = dict(
            cls_score=cls_score,
            bbox_pred=bbox_pred,
            bbox_feats=bbox_cls_feats)
        return bbox_results

最後に、対応するレジストリがそれらを見つけてロードできるように、ユーザーは mmdet/models/bbox_heads/__init__.pymmdet/models/roi_heads/__init__.py にモジュールを追加する必要があります。

または、ユーザーは、

custom_imports=dict(
    imports=['mmdet.models.roi_heads.double_roi_head', 'mmdet.models.roi_heads.bbox_heads.double_bbox_head'])

を設定ファイルに追加して、同じ目標を達成できます。

Double Head R-CNNの設定ファイルは次のとおりです。

_base_ = '../faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py'
model = dict(
    roi_head=dict(
        type='DoubleHeadRoIHead',
        reg_roi_scale_factor=1.3,
        bbox_head=dict(
            _delete_=True,
            type='DoubleConvFCBBoxHead',
            num_convs=4,
            num_fcs=2,
            in_channels=256,
            conv_out_channels=1024,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=80,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=2.0),
            loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=2.0))))

MMDetection 2.0以降、設定システムは設定を継承することをサポートしているため、ユーザーは変更に集中できます。Double Head R-CNNは、主に新しい DoubleHeadRoIHead と新しい DoubleConvFCBBoxHead を使用しており、引数は各モジュールの __init__ 関数に従って設定されます。

新しい損失の追加

境界ボックス回帰のために、MyLoss として新しい損失を追加すると仮定します。新しい損失関数を追加するには、ユーザーは mmdet/models/losses/my_loss.py でそれを実装する必要があります。デコレータ weighted_loss により、損失を各要素に対して重み付けできます。

import torch
import torch.nn as nn

from mmdet.registry import MODELS
from .utils import weighted_loss

@weighted_loss
def my_loss(pred, target):
    assert pred.size() == target.size() and target.numel() > 0
    loss = torch.abs(pred - target)
    return loss

@MODELS.register_module()
class MyLoss(nn.Module):

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(MyLoss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss_bbox = self.loss_weight * my_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss_bbox

次に、ユーザーはそれを mmdet/models/losses/__init__.py に追加する必要があります。

from .my_loss import MyLoss, my_loss

または、

custom_imports=dict(
    imports=['mmdet.models.losses.my_loss'])

を設定ファイルに追加して、同じ目標を達成できます。

を追加できます。それを使用するには、loss_xxx フィールドを変更します。MyLossは回帰用なので、ヘッドの loss_bbox フィールドを変更する必要があります。

loss_bbox=dict(type='MyLoss', loss_weight=1.0))