mmagic.models.editors.indexnet.indexnet_encoder

Module Contents

Classes

HolisticIndexBlock	Holistic Index Block.
DepthwiseIndexBlock	Depthwise index block.
InvertedResidual	Inverted residual layer for indexnet encoder.
IndexNetEncoder	Encoder for IndexNet.

Functions

build_index_block(in_channels, out_channels, kernel_size)

Build an conv block for IndexBlock.

mmagic.models.editors.indexnet.indexnet_encoder.build_index_block(in_channels, out_channels, kernel_size, stride=2, padding=0, groups=1, norm_cfg=dict(type='BN'), use_nonlinear=False, expansion=1)

Build an conv block for IndexBlock.

Parameters

• in_channels (int) – The input channels of the block.

• out_channels (int) – The output channels of the block.

• kernel_size (int) – The kernel size of the block.

• stride (int, optional) – The stride of the block. Defaults to 2.

• padding (int, optional) – The padding of the block. Defaults to 0.

• groups (int, optional) – The groups of the block. Defaults to 1.

• norm_cfg (dict, optional) – The norm config of the block. Defaults to dict(type=’BN’).

• use_nonlinear (bool, optional) – Whether use nonlinearity in the block. If true, a ConvModule with kernel size 1 will be appended and an ReLU6 nonlinearity will be added to the origin ConvModule. Defaults to False.

• expansion (int, optional) – Expansion ratio of the middle channels. Effective when use_nonlinear is true. Defaults to 1.

Returns

The built conv block.

Return type

nn.Module

class mmagic.models.editors.indexnet.indexnet_encoder.HolisticIndexBlock(in_channels, norm_cfg=dict(type='BN'), use_context=False, use_nonlinear=False)

Bases: mmengine.model.BaseModule

Holistic Index Block.

From https://arxiv.org/abs/1908.00672.

Parameters

• in_channels (int) – Input channels of the holistic index block.

• norm_cfg (dict) – Config dict for normalization layer. Default: dict(type=’BN’).

• use_context (bool, optional) – Whether use larger kernel size in index block. Refer to the paper for more information. Defaults to False.

• use_nonlinear (bool) – Whether add a non-linear conv layer in the index block. Default: False.

forward(x)

Forward function.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Encoder index feature and decoder index feature.

Return type

tuple(Tensor)

class mmagic.models.editors.indexnet.indexnet_encoder.DepthwiseIndexBlock(in_channels, norm_cfg=dict(type='BN'), use_context=False, use_nonlinear=False, mode='o2o')

Bases: mmengine.model.BaseModule

Depthwise index block.

From https://arxiv.org/abs/1908.00672.

Parameters

• in_channels (int) – Input channels of the holistic index block.

• norm_cfg (dict) – Config dict for normalization layer. Default: dict(type=’BN’).

• use_context (bool, optional) – Whether use larger kernel size in index block. Refer to the paper for more information. Defaults to False.

• use_nonlinear (bool) – Whether add a non-linear conv layer in the index blocks. Default: False.

• mode (str) – Mode of index block. Should be ‘o2o’ or ‘m2o’. In ‘o2o’ mode, the group of the conv layers is 1; In ‘m2o’ mode, the group of the conv layer is in_channels.

forward(x)

Forward function.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Encoder index feature and decoder index feature.

Return type

tuple(Tensor)

class mmagic.models.editors.indexnet.indexnet_encoder.InvertedResidual(in_channels, out_channels, stride, dilation, expand_ratio, norm_cfg, use_res_connect=False)

Bases: mmengine.model.BaseModule

Inverted residual layer for indexnet encoder.

It basically is a depthwise separable conv module. If expand_ratio is not one, then a conv module of kernel_size 1 will be inserted to change the input channels to in_channels * expand_ratio.

Parameters

• in_channels (int) – Input channels of the layer.

• out_channels (int) – Output channels of the layer.

• stride (int) – Stride of the depthwise separable conv module.

• dilation (int) – Dilation of the depthwise separable conv module.

• expand_ratio (float) – Expand ratio of the input channels of the depthwise separable conv module.

• norm_cfg (dict | None) – Config dict for normalization layer.

• use_res_connect (bool, optional) – Whether use shortcut connection. Defaults to False.

pad(inputs, kernel_size, dilation)

Pad input tensor.

Parameters

• inputs (Tensor) – Input tensor.

• kernel_size (int) – Kernel size of conv layer.

• dilation (int) – Dilation of conv layer.

Returns

Padded tensor

Return type

Tensor

forward(x)

Forward function.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Output feature map.

Return type

Tensor

class mmagic.models.editors.indexnet.indexnet_encoder.IndexNetEncoder(in_channels, out_stride=32, width_mult=1, index_mode='m2o', aspp=True, norm_cfg=dict(type='BN'), freeze_bn=False, use_nonlinear=True, use_context=True, init_cfg: Optional[dict] = None)

Bases: mmengine.model.BaseModule

Encoder for IndexNet.

Please refer to https://arxiv.org/abs/1908.00672.

Parameters

• in_channels (int, optional) – Input channels of the encoder.

• out_stride (int, optional) – Output stride of the encoder. For example, if out_stride is 32, the input feature map or image will be downsample to the 1/32 of original size. Defaults to 32.

• width_mult (int, optional) – Width multiplication factor of channel dimension in MobileNetV2. Defaults to 1.

• index_mode (str, optional) – Index mode of the index network. It must be one of {‘holistic’, ‘o2o’, ‘m2o’}. If it is set to ‘holistic’, then Holistic index network will be used as the index network. If it is set to ‘o2o’ (or ‘m2o’), when O2O (or M2O) Depthwise index network will be used as the index network. Defaults to ‘m2o’.

• aspp (bool, optional) – Whether use ASPP module to augment output feature. Defaults to True.

• norm_cfg (None | dict, optional) – Config dict for normalization layer. Defaults to dict(type=’BN’).

• freeze_bn (bool, optional) – Whether freeze batch norm layer. Defaults to False.

• use_nonlinear (bool, optional) – Whether use nonlinearity in index network. Refer to the paper for more information. Defaults to True.

• use_context (bool, optional) – Whether use larger kernel size in index network. Refer to the paper for more information. Defaults to True.

• init_cfg (dict, optional) – Initialization config dict. Default: None.

Raises

• ValueError – out_stride must 16 or 32.

• NameError – Supported index_mode are {‘holistic’, ‘o2o’, ‘m2o’}.

_make_layer(layer_setting, norm_cfg)

train(mode=True)

Set BatchNorm modules in the model to evaluation mode.

init_weights()

Init weights for the model.

Initialization is based on self._init_cfg

Parameters

pretrained (str, optional) – Path for pretrained weights. If given None, pretrained weights will not be loaded. Defaults to None.

forward(x)

Forward function.

Parameters

x (Tensor) – Input feature map with shape (N, C, H, W).

Returns

Output tensor, shortcut feature and decoder index feature.

Return type

dict