添加模块¶
在本教程中,我们将要介绍创建用户自定义模块的基本步骤。在学习创建自定义模块之前,建议先了解一下文件 models.md中模型的基本的概念。您可以自定义 models.md文件中涉及的所有模型组件,例如主干网络 (backbone), neck, head 和损失 (loss)。
添加新的主干网络¶
假如要创建新的主干网络 NewBackbone
。
新建一个文件
mmselfsup/models/backbones/new_backbone.py
并在其中实现新的主干网络NewBackbone
。
import torch.nn as nn
from mmselfsup.registry import MODELS
@MODELS.register_module()
class NewBackbone(nn.Module):
def __init__(self, *args, **kwargs):
pass
def forward(self, x): # should return a tuple
pass
def init_weights(self):
pass
def train(self, mode=True):
pass
导入新添加的主干网络到
mmselfsup/models/backbones/__init__.py
。
...
from .new_backbone import NewBackbone
__all__ = [
...,
'NewBackbone',
...
]
在配置文件中使用自定义的主干网络。
model = dict(
...
backbone=dict(
type='NewBackbone',
...),
...
)
添加新的 neck¶
您可以通过继承 mmengine 中的 BaseModule
来创建新的 neck,并且重定义其中的 forward
函数。我们在 mmengine 中有统一的权重初始化接口,您可以使用 init_cfg
来指定初始化函数和参数,或者可以重定义函数 init_weights
如果您希望使用自定义初始化方式。
所有的已有的 neck 都在 mmselfsup/models/necks
中。假如您要创建新的 neck NewNeck
。
新建一个文件
mmselfsup/models/necks/new_neck.py
并在其中实现NewNeck
。
from mmengine.model import BaseModule
from mmselfsup.registry import MODELS
@MODELS.register_module()
class NewNeck(BaseModule):
def __init__(self, *args, **kwargs):
super().__init__()
pass
def forward(self, x):
pass
您需要在 forward
函数中实现一些针对主干网络输出的操作,并将结果给 head。
导入新定义的 neck 模块到
mmselfsup/models/necks/__init__.py
。
...
from .new_neck import NewNeck
__all__ = [
...,
'NewNeck',
...
]
在配置文件中使用自定义的 neck 模块。
model = dict(
...
neck=dict(
type='NewNeck',
...),
...
)
添加新的 head¶
您可以通过继承 mmengine 中的 BaseModule
来创建新的 head,并且重定义其中的 forward
函数。
所有已有的 head 都在 mmselfsup/models/heads
文件中。假如您想创建新的head NewHead
。
创建文件
mmselfsup/models/heads/new_head.py
并在其中实现NewHead
。
from mmengine.model import BaseModule
from mmselfsup.registry import MODELS
@MODELS.register_module()
class NewHead(BaseModule):
def __init__(self, loss, **kwargs):
super().__init__()
# build loss
self.loss = MODELS.build(loss)
# other specific initializations
def forward(self, *args, **kwargs):
pass
您需要在 forward
函数中实现一些针对主干网络或 neck 输出的操作,并计算损失。请注意,损失模块应构建在 head 模块中以进行损失计算。
在
mmselfsup/models/heads/__init__.py
中导入新创建的 head 模块。
...
from .new_head import NewHead
__all__ = [
...,
'NewHead',
...
]
在配置文件中使用自定义的 head。
model = dict(
...
head=dict(
type='NewHead',
...),
...
)
添加新的损失¶
添加新的损失函数时,主要需要在 loss 模块中实现 forward
函数。同时您需要将 loss 模块也注册 (register) 为 MODELS
。
所有已有的损失函数都在 mmselfsup/models/losses
中。假如您想创建新的损失 NewLoss
。
创建文件
mmselfsup/models/losses/new_loss.py
并在其中实现NewLoss
。
from mmengine.model import BaseModule
from mmselfsup.registry import MODELS
@MODELS.register_module()
class NewLoss(BaseModule):
def __init__(self, *args, **kwargs):
super().__init__()
pass
def forward(self, *args, **kwargs):
pass
在
mmselfsup/models/losses/__init__.py
中导入新定义的 loss 模块。
...
from .new_loss import NewLoss
__all__ = [
...,
'NewLoss',
...
]
在配置文件中使用自定义的 loss 模块。
model = dict(
...
head=dict(
...
loss=dict(
type='NewLoss',
...),
...),
...
)
组合起来¶
在创建好上述的各个模块组件之后,我们需要创建一个新的算法 NewAlgorithm
来将各个组件按照逻辑顺序组合起来。NewAlgorithm
使用原始图像作为输入并输出损失函数值给优化器 (optimizer)。
创建文件
mmselfsup/models/algorithms/new_algorithm.py
并在其中实现NewAlgorithm
。
from mmselfsup.registry import MODELS
from .base import BaseModel
@MODELS.register_module()
class NewAlgorithm(BaseModel):
def __init__(self, backbone, neck=None, head=None, init_cfg=None):
super().__init__(init_cfg)
pass
def extract_feat(self, inputs, **kwargs):
pass
def loss(self, inputs, data_samples, **kwargs):
pass
def predict(self, inputs, data_samples, **kwargs):
pass
在
mmselfsup/models/algorithms/__init__.py
中导入新创建的算法模块NewAlgorithm
。
...
from .new_algorithm import NewAlgorithm
__all__ = [
...
'NewAlgorithm',
...
]
在配置文件中使用自定义的算法模块。
model = dict(
type='NewAlgorithm',
backbone=...,
neck=...,
head=...,
...
)