NonLinearNeck

class mmpretrain.models.necks.NonLinearNeck(in_channels, hid_channels, out_channels, num_layers=2, with_bias=False, with_last_bn=True, with_last_bn_affine=True, with_last_bias=False, with_avg_pool=True, norm_cfg={'type': 'SyncBN'}, init_cfg=[{'type': 'Constant', 'val': 1, 'layer': ['_BatchNorm', 'GroupNorm']}])

The non-linear neck.

Structure: fc-bn-[relu-fc-bn] where the substructure in [] can be repeated. For the default setting, the repeated time is 1. The neck can be used in many algorithms, e.g., SimCLR, BYOL, SimSiam.

Parameters:

• in_channels (int) – Number of input channels.

• hid_channels (int) – Number of hidden channels.

• out_channels (int) – Number of output channels.

• num_layers (int) – Number of fc layers. Defaults to 2.

• with_bias (bool) – Whether to use bias in fc layers (except for the last). Defaults to False.

• with_last_bn (bool) – Whether to add the last BN layer. Defaults to True.

• with_last_bn_affine (bool) – Whether to have learnable affine parameters in the last BN layer (set False for SimSiam). Defaults to True.

• with_last_bias (bool) – Whether to use bias in the last fc layer. Defaults to False.

• with_avg_pool (bool) – Whether to apply the global average pooling after backbone. Defaults to True.

• norm_cfg (dict) – Dictionary to construct and config norm layer. Defaults to dict(type=’SyncBN’).

• init_cfg (dict or list[dict], optional) – Initialization config dict.

forward(x)

Forward function.

Parameters:

x (Tuple[torch.Tensor]) – The feature map of backbone.

Returns:

The output features.

Return type:

Tuple[torch.Tensor]