Spaces:
Runtime error
Runtime error
File size: 4,533 Bytes
f549064 |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, Optional, Sequence, Tuple
import torch
import torch.nn as nn
from mmcv.cnn import build_activation_layer, build_norm_layer
from mmengine.model import BaseModule, ModuleList
from mmcls.registry import MODELS
from .cls_head import ClsHead
class LinearBlock(BaseModule):
"""Linear block for StackedLinearClsHead."""
def __init__(self,
in_channels,
out_channels,
dropout_rate=0.,
norm_cfg=None,
act_cfg=None,
init_cfg=None):
super().__init__(init_cfg=init_cfg)
self.fc = nn.Linear(in_channels, out_channels)
self.norm = None
self.act = None
self.dropout = None
if norm_cfg is not None:
self.norm = build_norm_layer(norm_cfg, out_channels)[1]
if act_cfg is not None:
self.act = build_activation_layer(act_cfg)
if dropout_rate > 0:
self.dropout = nn.Dropout(p=dropout_rate)
def forward(self, x):
"""The forward process."""
x = self.fc(x)
if self.norm is not None:
x = self.norm(x)
if self.act is not None:
x = self.act(x)
if self.dropout is not None:
x = self.dropout(x)
return x
@MODELS.register_module()
class StackedLinearClsHead(ClsHead):
"""Classifier head with several hidden fc layer and a output fc layer.
Args:
num_classes (int): Number of categories.
in_channels (int): Number of channels in the input feature map.
mid_channels (Sequence[int]): Number of channels in the hidden fc
layers.
dropout_rate (float): Dropout rate after each hidden fc layer,
except the last layer. Defaults to 0.
norm_cfg (dict, optional): Config dict of normalization layer after
each hidden fc layer, except the last layer. Defaults to None.
act_cfg (dict, optional): Config dict of activation function after each
hidden layer, except the last layer. Defaults to use "ReLU".
"""
def __init__(self,
num_classes: int,
in_channels: int,
mid_channels: Sequence[int],
dropout_rate: float = 0.,
norm_cfg: Optional[Dict] = None,
act_cfg: Optional[Dict] = dict(type='ReLU'),
**kwargs):
super(StackedLinearClsHead, self).__init__(**kwargs)
self.num_classes = num_classes
self.in_channels = in_channels
if self.num_classes <= 0:
raise ValueError(
f'num_classes={num_classes} must be a positive integer')
assert isinstance(mid_channels, Sequence), \
f'`mid_channels` of StackedLinearClsHead should be a sequence, ' \
f'instead of {type(mid_channels)}'
self.mid_channels = mid_channels
self.dropout_rate = dropout_rate
self.norm_cfg = norm_cfg
self.act_cfg = act_cfg
self._init_layers()
def _init_layers(self):
""""Init layers."""
self.layers = ModuleList()
in_channels = self.in_channels
for hidden_channels in self.mid_channels:
self.layers.append(
LinearBlock(
in_channels,
hidden_channels,
dropout_rate=self.dropout_rate,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg))
in_channels = hidden_channels
self.layers.append(
LinearBlock(
self.mid_channels[-1],
self.num_classes,
dropout_rate=0.,
norm_cfg=None,
act_cfg=None))
def pre_logits(self, feats: Tuple[torch.Tensor]) -> torch.Tensor:
"""The process before the final classification head.
The input ``feats`` is a tuple of tensor, and each tensor is the
feature of a backbone stage.
"""
x = feats[-1]
for layer in self.layers[:-1]:
x = layer(x)
return x
@property
def fc(self):
"""Full connected layer."""
return self.layers[-1]
def forward(self, feats: Tuple[torch.Tensor]) -> torch.Tensor:
"""The forward process."""
pre_logits = self.pre_logits(feats)
# The final classification head.
cls_score = self.fc(pre_logits)
return cls_score
|