Add New Modules — MMSegmentation 1.2.2 documentation (original) (raw)

Develop new components¶

We can customize all the components introduced at the model documentation, such as backbone, head, loss function and data preprocessor.

Add new backbones¶

Here we show how to develop a new backbone with an example of MobileNet.

Create a new file mmseg/models/backbones/mobilenet.py.
import torch.nn as nn
from mmseg.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
def init_weights(self, pretrained=None):
pass
Import the module in mmseg/models/backbones/__init__.py.
from .mobilenet import MobileNet
Use it in your config file.
model = dict(
...
backbone=dict(
type='MobileNet',
arg1=xxx,
arg2=xxx),
...

Add new heads¶

In MMSegmentation, we provide a BaseDecodeHead for developing all segmentation heads. All newly implemented decode heads should be derived from it. Here we show how to develop a new head with the example of PSPNet as the following.

First, add a new decode head in mmseg/models/decode_heads/psp_head.py. PSPNet implements a decode head for segmentation decode. To implement a decode head, we need to implement three functions of the new module as the following.

from mmseg.registry import MODELS

@MODELS.register_module() class PSPHead(BaseDecodeHead):

def __init__(self, pool_scales=(1, 2, 3, 6), **kwargs):
    super(PSPHead, self).__init__(**kwargs)

def init_weights(self):
    pass

def forward(self, inputs):
    pass

Next, the users need to add the module in the mmseg/models/decode_heads/__init__.py, thus the corresponding registry could find and load them.

To config file of PSPNet is as the following

norm_cfg = dict(type='SyncBN', requires_grad=True) model = dict( type='EncoderDecoder', pretrained='pretrain_model/resnet50_v1c_trick-2cccc1ad.pth', backbone=dict( type='ResNetV1c', depth=50, num_stages=4, out_indices=(0, 1, 2, 3), dilations=(1, 1, 2, 4), strides=(1, 2, 1, 1), norm_cfg=norm_cfg, norm_eval=False, style='pytorch', contract_dilation=True), decode_head=dict( type='PSPHead', in_channels=2048, in_index=3, channels=512, pool_scales=(1, 2, 3, 6), dropout_ratio=0.1, num_classes=19, norm_cfg=norm_cfg, align_corners=False, loss_decode=dict( type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)))

Add new loss¶

Assume you want to add a new loss as MyLoss for segmentation decode. To add a new loss function, the users need to implement it in mmseg/models/losses/my_loss.py. The decorator weighted_loss enables the loss to be weighted for each element.

import torch import torch.nn as nn

from mmseg.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 = self.loss_weight * my_loss(
        pred, target, weight, reduction=reduction, avg_factor=avg_factor)
    return loss

Then the users need to add it in the mmseg/models/losses/__init__.py.

from .my_loss import MyLoss, my_loss

To use it, modify the loss_xxx field. Then you need to modify the loss_decode field in the head.loss_weight could be used to balance multiple losses.

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

Add new data preprocessor¶

In MMSegmentation 1.x versions, we use SegDataPreProcessor to copy data to the target device and preprocess the data into the model input format as default. Here we show how to develop a new data preprocessor.

Create a new file mmseg/models/my_datapreprocessor.py.
from mmengine.model import BaseDataPreprocessor
from mmseg.registry import MODELS
@MODELS.register_module()
class MyDataPreProcessor(BaseDataPreprocessor):
def init(self, **kwargs):
super().init(**kwargs)
def forward(self, data: dict, training: bool=False) -> Dict[str, Any]:
# TODO Define the logic for data pre-processing in the forward method
pass
Import your data preprocessor in mmseg/models/__init__.py
from .my_datapreprocessor import MyDataPreProcessor
Use it in your config file.
model = dict(
data_preprocessor=dict(type='MyDataPreProcessor)
...
)

Develop new segmentors¶

The segmentor is an algorithmic architecture in which users can customize their algorithms by adding customized components and defining the logic of algorithm execution. Please refer to the model document for more details.

Since the BaseSegmentor in MMSegmentation unifies three modes for a forward process, to develop a new segmentor, users need to overwrite loss, predict and _forward methods corresponding to the loss, predict and tensor modes.

Here we show how to develop a new segmentor.

Create a new file mmseg/models/segmentors/my_segmentor.py.
from typing import Dict, Optional, Union
import torch
from mmseg.registry import MODELS
from mmseg.models import BaseSegmentor
@MODELS.register_module()
class MySegmentor(BaseSegmentor):
def init(self, **kwargs):
super().init(**kwargs)
# TODO users should build components of the network here
def loss(self, inputs: Tensor, data_samples: SampleList) -> dict:
"""Calculate losses from a batch of inputs and data samples."""
pass
def predict(self, inputs: Tensor, data_samples: OptSampleList=None) -> SampleList:
"""Predict results from a batch of inputs and data samples with post-
processing."""
pass
def _forward(self,
inputs: Tensor,
data_samples: OptSampleList = None) -> Tuple[List[Tensor]]:
"""Network forward process.
Usually includes backbone, neck and head forward without any post-
processing.
"""
pass
Import your segmentor in mmseg/models/segmentors/__init__.py.
from .my_segmentor import MySegmentor
Use it in your config file.
model = dict(
type='MySegmentor'
...
)