教程 4: 自定义模型¶
我们简单地把模型的各个组件分为五类：
主干网络 (backbone)：通常是一个用来提取特征图 (feature map) 的全卷积网络 (FCN network)，例如：ResNet, MobileNet。
Neck：主干网络和 Head 之间的连接部分，例如：FPN, PAFPN。
Head：用于具体任务的组件，例如：边界框预测和掩码预测。
区域提取器 (roi extractor)：从特征图中提取 RoI 特征，例如：RoI Align。
损失 (loss)：在 Head 组件中用于计算损失的部分，例如：FocalLoss, L1Loss, GHMLoss.
开发新的组件¶
添加一个新的主干网络¶
这里，我们以 MobileNet 为例来展示如何开发新组件。
1. 定义一个新的主干网络（以 MobileNet 为例）¶
新建一个文件 mmdet/models/backbones/mobilenet.py
import torch.nn as nn
from ..builder import BACKBONES
@BACKBONES.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),
...
添加新的 Neck¶
1. 定义一个 Neck（以 PAFPN 为例）¶
新建一个文件 mmdet/models/necks/pafpn.py
from ..builder import NECKS
@NECKS.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.py'],
allow_failed_imports=False)
到配置文件以避免原始代码被修改。
3. 修改配置文件¶
neck=dict(
type='PAFPN',
in_channels=[256, 512, 1024, 2048],
out_channels=256,
num_outs=5)
添加新的 Head¶
我们以 Double Head R-CNN 为例来展示如何添加一个新的 Head。
首先，添加一个新的 bbox head 到 mmdet/models/roi_heads/bbox_heads/double_bbox_head.py。
Double Head R-CNN 在目标检测上实现了一个新的 bbox head。为了实现 bbox head，我们需要使用如下的新模块中三个函数。
from mmdet.models.builder import HEADS
from .bbox_head import BBoxHead
@HEADS.register_module()
class DoubleConvFCBBoxHead(BBoxHead):
r"""Bbox head used in Double-Head R-CNN
/-> cls
/-> shared convs ->
\-> reg
roi features
/-> cls
\-> shared fc ->
\-> reg
""" # noqa: W605
def __init__(self,
num_convs=0,
num_fcs=0,
conv_out_channels=1024,
fc_out_channels=1024,
conv_cfg=None,
norm_cfg=dict(type='BN'),
**kwargs):
kwargs.setdefault('with_avg_pool', True)
super(DoubleConvFCBBoxHead, self).__init__(**kwargs)
def forward(self, x_cls, x_reg):
然后，如有必要，实现一个新的 bbox head。我们打算从 StandardRoIHead 来继承新的 DoubleHeadRoIHead。我们可以发现 StandardRoIHead 已经实现了下述函数。
import torch
from mmdet.core import bbox2result, bbox2roi, build_assigner, build_sampler
from ..builder import HEADS, build_head, build_roi_extractor
from .base_roi_head import BaseRoIHead
from .test_mixins import BBoxTestMixin, MaskTestMixin
@HEADS.register_module()
class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
"""Simplest base roi head including one bbox head and one mask head.
"""
def init_assigner_sampler(self):
def init_bbox_head(self, bbox_roi_extractor, bbox_head):
def init_mask_head(self, mask_roi_extractor, mask_head):
def forward_dummy(self, x, proposals):
def forward_train(self,
x,
img_metas,
proposal_list,
gt_bboxes,
gt_labels,
gt_bboxes_ignore=None,
gt_masks=None):
def _bbox_forward(self, x, rois):
def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels,
img_metas):
def _mask_forward_train(self, x, sampling_results, bbox_feats, gt_masks,
img_metas):
def _mask_forward(self, x, rois=None, pos_inds=None, bbox_feats=None):
def simple_test(self,
x,
proposal_list,
img_metas,
proposals=None,
rescale=False):
"""Test without augmentation."""
Double Head 的修改主要在 bbox_forward 的逻辑中，且它从 StandardRoIHead 中继承了其他逻辑。在 mmdet/models/roi_heads/double_roi_head.py 中，我们用下述代码实现新的 bbox head：
from ..builder import HEADS
from .standard_roi_head import StandardRoIHead
@HEADS.register_module()
class DoubleHeadRoIHead(StandardRoIHead):
"""RoI head for Double Head RCNN
https://arxiv.org/abs/1904.06493
"""
def __init__(self, reg_roi_scale_factor, **kwargs):
super(DoubleHeadRoIHead, self).__init__(**kwargs)
self.reg_roi_scale_factor = reg_roi_scale_factor
def _bbox_forward(self, x, rois):
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__.py 和 mmdet/models/roi_heads/__init__.py 以使相关的注册表可以找到并加载他们。
或者，用户可以添加：
custom_imports=dict(
imports=['mmdet.models.roi_heads.double_roi_head', 'mmdet.models.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 ..builder import LOSSES
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
@LOSSES.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 是用于回归的，你需要在 Head 中修改 loss_xxx 字段。
loss_bbox=dict(type='MyLoss', loss_weight=1.0))