| |
| import torch |
| import torch.nn as nn |
| from mmengine.model import BaseModule |
|
|
| from mmdet.registry import MODELS |
|
|
|
|
| @MODELS.register_module() |
| class GlobalAveragePooling(BaseModule): |
| """Global Average Pooling neck. |
| |
| Note that we use `view` to remove extra channel after pooling. We do not |
| use `squeeze` as it will also remove the batch dimension when the tensor |
| has a batch dimension of size 1, which can lead to unexpected errors. |
| """ |
|
|
| def __init__(self, kernel_size=None, stride=None): |
| super(GlobalAveragePooling, self).__init__() |
| if kernel_size is None and stride is None: |
| self.gap = nn.AdaptiveAvgPool2d((1, 1)) |
| else: |
| self.gap = nn.AvgPool2d(kernel_size, stride) |
|
|
| def forward(self, inputs): |
| if isinstance(inputs, tuple): |
| outs = tuple([self.gap(x) for x in inputs]) |
| outs = tuple([ |
| out.view(x.size(0), |
| torch.tensor(out.size()[1:]).prod()) |
| for out, x in zip(outs, inputs) |
| ]) |
| elif isinstance(inputs, torch.Tensor): |
| outs = self.gap(inputs) |
| outs = outs.view( |
| inputs.size(0), |
| torch.tensor(outs.size()[1:]).prod()) |
| else: |
| raise TypeError('neck inputs should be tuple or torch.tensor') |
| return outs |
|
|
