mirrors
/
mmselfsup
mirror of https://github.com/open-mmlab/mmselfsup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[Docs] refactor add_modules.md (#524) 

* [Docs] refactor add_modules.md

* [Docs] fix lint
pull/582/head
Jiahao Xie 2022-10-24 15:23:42 +08:00 committed by Yixiao Fang
parent 5a0e0eedc4
commit d142699265
1 changed files with 160 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

249
docs/en/advanced_guides/add_modules.md
View File

@ -1,204 +1,275 @@
# Add Modules # Add Modules
In this tutorial, we introduce the basic steps to create your customized modules. Before learning to create your customized modules, it is recommended to learn the basic concept of models in file [models.md](models.md). You can customize all the components introduced in [models.md](models.md), such as **backbone**, **neck**, **head** and **loss**.
- [Add Modules](#add-modules) - [Add Modules](#add-modules)
- [Add new backbone](#add-new-backbone) - [Add a new backbone](#add-a-new-backbone)
- [Add new necks](#add-new-necks) - [Add a new neck](#add-a-new-neck)
- [Add new loss](#add-new-loss) - [Add a new head](#add-a-new-head)
- [Add a new loss](#add-a-new-loss)
- [Combine all](#combine-all) - [Combine all](#combine-all)
In self-supervised learning domain, each model can be divided into following four parts: ## Add a new backbone
- backbone: used to extract image's feature Assume you are going to create a new backbone `NewBackbone`.
- projection head: projects feature extracted by backbone to another space
- loss: loss function the model will optimize
- memory bank(optional): some methods, `e.g. odc`, need extract memory bank to store image's feature.
## Add new backbone 1. Create a new file `mmselfsup/models/backbones/new_backbone.py` and implement `NewBackbone` in it.
Assuming we are going to create a customized backbone `CustomizedBackbone`
1.Create a new file `mmselfsup/models/backbones/customized_backbone.py` and implement `CustomizedBackbone` in it.
```python ```python
import torch.nn as nn import torch.nn as nn
from ..builder import BACKBONES
@BACKBONES.register_module() from mmselfsup.registry import MODELS
class CustomizedBackbone(nn.Module):
def __init__(self, **kwargs):
## TODO @MODELS.register_module()
class NewBackbone(nn.Module):
def forward(self, x): def __init__(self, *args, **kwargs):
pass
## TODO def forward(self, x): # should return a tuple
pass
def init_weights(self, pretrained=None): def init_weights(self):
pass
## TODO
def train(self, mode=True): def train(self, mode=True):
pass
## TODO
``` ```
2.Import the customized backbone in `mmselfsup/models/backbones/__init__.py`. 2. Import the new backbone module in `mmselfsup/models/backbones/__init__.py`.
```python ```python
from .customized_backbone import CustomizedBackbone ...
from .new_backbone import NewBackbone
__all__ = [ __all__ = [
..., 'CustomizedBackbone' ...,
'NewBackbone',
...
] ]
``` ```
3.Use it in your config file. 3. Use it in your config file.
```python ```python
model = dict( model = dict(
... ...
backbone=dict( backbone=dict(
type='CustomizedBackbone', type='NewBackbone',
...), ...),
... ...
) )
``` ```
## Add new necks ## Add a new neck
we include all projection heads in `mmselfsup/models/necks`. Assuming we are going to create a `CustomizedProjHead`. You can write a new neck inherited from `BaseModule` from mmengine, and overwrite `forward`. We have a unified interface for weight initialization in mmengine, you can use `init_cfg` to specify the initialization function and arguments, or overwrite `init_weights` if you prefer customized initialization.
1.Create a new file `mmselfsup/models/necks/customized_proj_head.py` and implement `CustomizedProjHead` in it. We include all necks in `mmselfsup/models/necks`. Assume you are going to create a new neck `NewNeck`.
1. Create a new file `mmselfsup/models/necks/new_neck.py` and implement `NewNeck` in it.
```python ```python
import torch.nn as nn from mmengine.model import BaseModule
from mmcv.runner import BaseModule
from ..builder import NECKS from mmselfsup.registry import MODELS
@NECKS.register_module() @MODELS.register_module()
class CustomizedProjHead(BaseModule): class NewNeck(BaseModule):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super(CustomizedProjHead, self).__init__(init_cfg) super().__init__()
## TODO pass
def forward(self, x): def forward(self, x):
## TODO pass
``` ```
You need to implement the forward function, which takes the feature from the backbone and outputs the projected feature. You need to implement the `forward` function, which applies some operations on the output from the backbone and forwards the results to the head.
2.Import the `CustomizedProjHead` in `mmselfsup/models/necks/__init__`. 2. Import the new neck module in `mmselfsup/models/necks/__init__.py`.
```python ```python
from .customized_proj_head import CustomizedProjHead ...
from .new_neck import NewNeck
__all__ = [ __all__ = [
..., ...,
CustomizedProjHead, 'NewNeck',
... ...
] ]
``` ```
3.Use it in your config file. 3. Use it in your config file.
```python ```python
model = dict( model = dict(
..., ...
neck=dict( neck=dict(
type='CustomizedProjHead', type='NewNeck',
...), ...),
...) ...
)
``` ```
## Add new loss ## Add a new head
To add a new loss function, we mainly implement the `forward` function in the loss module. You can write a new head inherited from `BaseModule` from mmengine, and overwrite `forward`.
1.Create a new file `mmselfsup/models/heads/customized_head.py` and implement your customized `CustomizedHead` in it. We include all heads in `mmselfsup/models/heads`. Assume you are going to create a new head `NewHead`.
1. Create a new file `mmselfsup/models/heads/new_head.py` and implement `NewHead` in it.
```python ```python
import torch from mmengine.model import BaseModule
import torch.nn as nn
from mmcv.runner import BaseModule
from ..builder import HEADS from mmselfsup.registry import MODELS
@HEADS.register_module() @MODELS.register_module()
class CustomizedHead(BaseModule): class NewHead(BaseModule):
def __init__(self, *args, **kwargs): def __init__(self, loss, **kwargs):
super(CustomizedHead, self).__init__() super().__init__()
# build loss
## TODO self.loss = MODELS.build(loss)
# other specific initializations
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
pass
## TODO
``` ```
2.Import the module in `mmselfsup/models/heads/__init__.py` You need to implement the `forward` function, which applies some operations on the output from the neck/backbone and computes the loss. Please note that the loss module should be built in the head module for the loss computation.
2. Import the new head module in `mmselfsup/models/heads/__init__.py`.
```python ```python
from .customized_head import CustomizedHead ...
from .new_head import NewHead
__all__ = [..., CustomizedHead, ...] __all__ = [
...,
'NewHead',
...
]
``` ```
3.Use it in your config file. 3. Use it in your config file.
```python ```python
model = dict( model = dict(
...
head=dict(
type='NewHead',
...),
...
)
```
## Add a new loss
To add a new loss function, we mainly implement the `forward` function in the loss module. We should register the loss module as `MODELS` as well.
We include all losses in `mmselfsup/models/losses`. Assume you are going to create a new loss `NewLoss`.
1. Create a new file `mmselfsup/models/losses/new_loss.py` and implement `NewLoss` in it.
```python
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
```
2. Import the new loss module in `mmselfsup/models/losses/__init__.py`
```python
...
from .new_loss import NewLoss
__all__ = [
..., ...,
head=dict(type='CustomizedHead') 'NewLoss',
) ...
]
```
3. Use it in your config file.
```python
model = dict(
...
head=dict(
...
loss=dict(
type='NewLoss',
...),
...),
...
)
``` ```
## Combine all ## Combine all
After creating each component, mentioned above, we need to create a `CustomizedAlgorithm` to organize them logically. And the `CustomizedAlgorithm` takes raw images as inputs and outputs the loss to the optimizer. After creating each component mentioned above, we need to create a new algorithm `NewAlgorithm` to organize them logically. `NewAlgorithm` takes raw images as inputs and outputs the loss to the optimizer.
1.Create a new file `mmselfsup/models/algorithms/customized_algorithm.py` and implement `CustomizedAlgorithm` in it. 1. Create a new file `mmselfsup/models/algorithms/new_algorithm.py` and implement `NewAlgorithm` in it.
```python ```python
# Copyright (c) OpenMMLab. All rights reserved. from mmselfsup.registry import MODELS
import torch
from ..builder import ALGORITHMS, build_backbone, build_head, build_neck
from ..utils import GatherLayer
from .base import BaseModel from .base import BaseModel
@ALGORITHMS.register_module() @MODELS.register_module()
class CustomizedAlgorithm(BaseModel): class NewAlgorithm(BaseModel):
def __init__(self, backbone, neck=None, head=None, init_cfg=None): def __init__(self, backbone, neck=None, head=None, init_cfg=None):
super(SimCLR, self).__init__(init_cfg) super().__init__(init_cfg)
pass
## TODO def extract_feat(self, inputs, **kwargs):
pass
def forward_train(self, img, **kwargs): def loss(self, inputs, data_samples, **kwargs):
pass
## TODO def predict(self, inputs, data_samples, **kwargs):
pass
``` ```
2.Import the module in `mmselfsup/models/algorithms/__init__.py` 2. Import the new algorithm module in `mmselfsup/models/algorithms/__init__.py`
```python ```python
from .customized_algorithm import CustomizedAlgorithm ...
from .new_algorithm import NewAlgorithm
__all__ = [..., CustomizedAlgorithm, ...] __all__ = [
...,
'NewAlgorithm',
...
]
``` ```
3.Use it in your config file. 3. Use it in your config file.
```python ```python
model = dict( model = dict(
type='CustomizedAlgorightm', type='NewAlgorithm',
backbone=..., backbone=...,
neck=..., neck=...,
head=...) head=...,
...
)
``` ```