Shortcuts

注意

您正在阅读 MMSelfSup 0.x 版本的文档,而 MMSelfSup 0.x 版本将会在 2022 年末 开始逐步停止维护。我们建议您及时升级到 MMSelfSup 1.0.0rc 版本,享受由 OpenMMLab 2.0 带来的更多新特性和更佳的性能表现。阅读 MMSelfSup 1.0.0rc 的 发版日志, 代码文档 获取更多信息。

教程 3:添加新的模块

在自监督学习领域,每个模型可以被分为以下四个部分:

  • backbone:用于提取图像特征。

  • projection head:将 backbone 提取的特征映射到另一空间。

  • loss:用于模型优化的损失函数。

  • memory bank(可选):一些方法(例如 odc ),需要额外的 memory bank 用于存储图像特征。

添加新的 backbone

假设我们要创建一个自定义的 backbone CustomizedBackbone

1.创建新文件 mmselfsup/models/backbones/customized_backbone.py 并在其中实现 CustomizedBackbone

import torch.nn as nn
from ..builder import BACKBONES

@BACKBONES.register_module()
class CustomizedBackbone(nn.Module):

    def __init__(self, **kwargs):

        ## TODO

    def forward(self, x):

        ## TODO

    def init_weights(self, pretrained=None):

        ## TODO

    def train(self, mode=True):

        ## TODO

2.在 mmselfsup/models/backbones/__init__.py 中导入自定义的 backbone。

from .customized_backbone import CustomizedBackbone

__all__ = [
    ..., 'CustomizedBackbone'
]

3.在你的配置文件中使用它。

model = dict(
    ...
    backbone=dict(
        type='CustomizedBackbone',
        ...),
    ...
)

添加新的 Necks

我们在 mmselfsup/models/necks 中包含了所有的 projection heads。假设我们要创建一个 CustomizedProjHead

1.创建一个新文件 mmselfsup/models/necks/customized_proj_head.py 并在其中实现 CustomizedProjHead

import torch.nn as nn
from mmcv.runner import BaseModule

from ..builder import NECKS


@NECKS.register_module()
class CustomizedProjHead(BaseModule):

    def __init__(self, *args, **kwargs):
        super(CustomizedProjHead, self).__init__(init_cfg)
        ## TODO
    def forward(self, x):
        ## TODO

你需要实现前向函数,该函数从 backbone 中获取特征,并输出映射后的特征。

2.在 mmselfsup/models/necks/__init__ 中导入 CustomizedProjHead

from .customized_proj_head import CustomizedProjHead

__all__ = [
    ...,
    CustomizedProjHead,
    ...
]

3.在你的配置文件中使用它。

model = dict(
    ...,
    neck=dict(
        type='CustomizedProjHead',
        ...),
   ...)

添加新的损失

为了增加一个新的损失函数,我们主要在损失模块中实现 forward 函数。

1.创建一个新的文件 mmselfsup/models/heads/customized_head.py 并在其中实现你自定义的 CustomizedHead

import torch
import torch.nn as nn
from mmcv.runner import BaseModule

from ..builder import HEADS


@HEADS.register_module()
class CustomizedHead(BaseModule):

    def __init__(self, *args, **kwargs):
        super(CustomizedHead, self).__init__()

        ## TODO

    def forward(self, *args, **kwargs):

        ## TODO

2.在 mmselfsup/models/heads/__init__.py 中导入该模块。

from .customized_head import CustomizedHead

__all__ = [..., CustomizedHead, ...]

3.在你的配置文件中使用它。

model = dict(
    ...,
    head=dict(type='CustomizedHead')
    )

合并所有改动

在创建了上述每个组件后,我们需要创建一个 CustomizedAlgorithm 来有逻辑的将他们组织到一起。 CustomizedAlgorithm 接收原始图像作为输入,并将损失输出给优化器。

1.创建一个新文件 mmselfsup/models/algorithms/customized_algorithm.py 并在其中实现 CustomizedAlgorithm

# Copyright (c) OpenMMLab. All rights reserved.
import torch

from ..builder import ALGORITHMS, build_backbone, build_head, build_neck
from ..utils import GatherLayer
from .base import BaseModel


@ALGORITHMS.register_module()
class CustomizedAlgorithm(BaseModel):

    def __init__(self, backbone, neck=None, head=None, init_cfg=None):
        super(SimCLR, self).__init__(init_cfg)

        ## TODO

    def forward_train(self, img, **kwargs):

        ## TODO

2.在 mmselfsup/models/algorithms/__init__.py 中导入该模块。

from .customized_algorithm import CustomizedAlgorithm

__all__ = [..., CustomizedAlgorithm, ...]

3.在你的配置文件中使用它。

model = dict(
    type='CustomizedAlgorightm',
    backbone=...,
    neck=...,
    head=...)
Read the Docs v: latest
Versions
latest
stable
1.x
dev-1.x
dev
Downloads
pdf
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.