Spaces:
Runtime error
Runtime error
File size: 3,100 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
from mmcv.cnn import build_activation_layer, build_norm_layer
from mmengine.model import BaseModule
from mmpretrain.registry import MODELS
@MODELS.register_module()
class LinearNeck(BaseModule):
"""Linear neck with Dimension projection.
Args:
in_channels (int): Number of channels in the input.
out_channels (int): Number of channels in the output.
gap_dim (int): Dimensions of each sample channel, can be one of
{0, 1, 2, 3}. Defaults to 0.
norm_cfg (dict, optional): dictionary to construct and
config norm layer. Defaults to dict(type='BN1d').
act_cfg (dict, optional): dictionary to construct and
config activate layer. Defaults to None.
init_cfg (dict, optional): dictionary to initialize weights.
Defaults to None.
"""
def __init__(self,
in_channels: int,
out_channels: int,
gap_dim: int = 0,
norm_cfg: Optional[dict] = dict(type='BN1d'),
act_cfg: Optional[dict] = None,
init_cfg: Optional[dict] = None):
super().__init__(init_cfg=init_cfg)
self.in_channels = in_channels
self.out_channels = out_channels
self.norm_cfg = copy.deepcopy(norm_cfg)
self.act_cfg = copy.deepcopy(act_cfg)
assert gap_dim in [0, 1, 2, 3], 'GlobalAveragePooling dim only ' \
f'support {0, 1, 2, 3}, get {gap_dim} instead.'
if gap_dim == 0:
self.gap = nn.Identity()
elif gap_dim == 1:
self.gap = nn.AdaptiveAvgPool1d(1)
elif gap_dim == 2:
self.gap = nn.AdaptiveAvgPool2d((1, 1))
elif gap_dim == 3:
self.gap = nn.AdaptiveAvgPool3d((1, 1, 1))
self.fc = nn.Linear(in_features=in_channels, out_features=out_channels)
if norm_cfg:
self.norm = build_norm_layer(norm_cfg, out_channels)[1]
else:
self.norm = nn.Identity()
if act_cfg:
self.act = build_activation_layer(act_cfg)
else:
self.act = nn.Identity()
def forward(self, inputs: Union[Tuple,
torch.Tensor]) -> Tuple[torch.Tensor]:
"""forward function.
Args:
inputs (Union[Tuple, torch.Tensor]): The features extracted from
the backbone. Multiple stage inputs are acceptable but only
the last stage will be used.
Returns:
Tuple[torch.Tensor]: A tuple of output features.
"""
assert isinstance(inputs, (tuple, torch.Tensor)), (
'The inputs of `LinearNeck` must be tuple or `torch.Tensor`, '
f'but get {type(inputs)}.')
if isinstance(inputs, tuple):
inputs = inputs[-1]
x = self.gap(inputs)
x = x.view(x.size(0), -1)
out = self.act(self.norm(self.fc(x)))
return (out, )
|