目录
1、定制优化器
1.1、优化器注册类
1.2、优化器定义
2、开发新组件
2.1、添加新的主干(backbone)
2.2、添加新的头(head)
2.3、增加新的损失
1、定制优化器
1.1、优化器注册类
mmsegmentation中未对优化器进行封装，真实的封装及注册类在mmcv中完成。我们可以在schedule默认配置中看到优化器的字典型配置。
在mmseg/api/train.py中，由build_optimizer完成构造。
进一步代码跟进，在mmcv中实现该注册类。
所以，不用纠结，工程代码中为什么字典型优化器配置，如何转为pytorch的优化器类。
1.2、优化器定义
假设定义一个名为优化MyOptimizer，其中有参数a，b和c。需要首先在文件中实现新的优化器，例如，在mmseg/core/optimizer/my_optimizer.py：
from mmcv.runner import OPTIMIZERS
from torch.optim import Optimizer
@OPTIMIZERS.register_module
class MyOptimizer(Optimizer):
def __init__(self, a, b, c):
pass然后在其中添加此模块，mmseg/core/optimizer/__init__.py以便注册表将找到新模块并将其添加：
from .my_optimizer import MyOptimizer然后，您可以MyOptimizer在optimizer配置文件的字段中使用。在配置中，优化器由字段定义，optimizer如下所示：
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)要使用自己的优化器，可以将字段更改为：
optimizer = dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value)我们已经支持使用由PyTorch实现的所有优化器，唯一的修改就是更改optimizer配置文件的字段。例如，如果您要使用ADAM，尽管性能会下降很多，但可以进行如下修改。
optimizer = dict(type='Adam', lr=0.0003, weight_decay=0.0001)用户可以按照PyTorch的API文档直接设置参数。
2、开发新组件
MMSegmentation中主要有2种类型的组件。
1）backbone：通常是卷积网络堆栈以提取特征图，例如ResNet，HRNet。
2）head：用于语义分割图解码的组件。
2.1、添加新的主干(backbone)
在这里，我们以MobileNet为例说明如何开发新组件。完整实现方式可以参考mmseg/models/backbone/mobilenet_v2.py。
1）创建一个新文件mmseg/models/backbones/mobilenet.py。以下三个类函数是必须实现的。
import torch.nn as nn
from ..registry import BACKBONES
@BACKBONES.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):
pass2）在中导入模块mmseg/models/backbones/__init__.py。
from .mobilenet import MobileNet3）在配置文件中使用它。
model = dict( ...
backbone=dict( type='MobileNet', arg1=xxx, arg2=xxx), ...2.2、添加新的头(head)
在MMSegmentation中，为所有head提供了基类BaseDecodeHead。所有新实现的head都应从中派生。在这里，我们以PSPNet为例说明如何开发新的头，如下所示。
首先，在中添加一个新的head：mmseg/models/decode_heads/psp_head.py。PSPNet实现了用于解码的分割head。要实现head，基本上我们需要实现以下新模块的三个功能。完整参考：psp_head.py中PSPHead中完整实现。
@HEADS.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接下来，用户需要在其中添加模块，mmseg/models/decode_heads/__init__.py因此相应的注册表可以找到并加载它们。
如下配置PSPNet文件：
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)
))2.3、增加新的损失
假设要添加新的损失MyLoss。要添加新的损失函数，用户需要在mmseg/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 = self.loss_weight * my_loss(
pred, target, weight, reduction=reduction, avg_factor=avg_factor)
return loss然后，用户需要将其添加到中mmseg/models/losses/__init__.py。
from .my_loss import MyLoss, my_loss要使用它，修改loss_xxx字段。然后，您需要修改head的loss_decode字段。 loss_weight可以用来平衡多重损失。
loss_decode=dict(type='MyLoss', loss_weight=1.0))
