Spaces:
Runtime error
Runtime error
File size: 2,637 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
from mmengine.model import BaseModule
from mmpretrain.registry import MODELS
@MODELS.register_module()
class DenseCLNeck(BaseModule):
"""The non-linear neck of DenseCL.
Single and dense neck in parallel: fc-relu-fc, conv-relu-conv.
Borrowed from the authors' `code <https://github.com/WXinlong/DenseCL>`_.
Args:
in_channels (int): Number of input channels.
hid_channels (int): Number of hidden channels.
out_channels (int): Number of output channels.
num_grid (int): The grid size of dense features. Defaults to None.
init_cfg (dict or list[dict], optional): Initialization config dict.
Defaults to None.
"""
def __init__(self,
in_channels: int,
hid_channels: int,
out_channels: int,
num_grid: Optional[int] = None,
init_cfg: Optional[Union[dict, List[dict]]] = None) -> None:
super().__init__(init_cfg)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.mlp = nn.Sequential(
nn.Linear(in_channels, hid_channels), nn.ReLU(inplace=True),
nn.Linear(hid_channels, out_channels))
self.with_pool = True if num_grid is not None else False
if self.with_pool:
self.pool = nn.AdaptiveAvgPool2d((num_grid, num_grid))
self.mlp2 = nn.Sequential(
nn.Conv2d(in_channels, hid_channels, 1), nn.ReLU(inplace=True),
nn.Conv2d(hid_channels, out_channels, 1))
self.avgpool2 = nn.AdaptiveAvgPool2d((1, 1))
def forward(self, x: Tuple[torch.Tensor]) -> Tuple[torch.Tensor]:
"""Forward function of neck.
Args:
x (Tuple[torch.Tensor]): feature map of backbone.
Returns:
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- ``avgpooled_x``: Global feature vectors.
- ``x``: Dense feature vectors.
- ``avgpooled_x2``: Dense feature vectors for queue.
"""
assert len(x) == 1
x = x[0]
avgpooled_x = self.avgpool(x)
avgpooled_x = self.mlp(avgpooled_x.view(avgpooled_x.size(0), -1))
if self.with_pool:
x = self.pool(x) # sxs
x = self.mlp2(x) # sxs: bxdxsxs
avgpooled_x2 = self.avgpool2(x) # 1x1: bxdx1x1
x = x.view(x.size(0), x.size(1), -1) # bxdxs^2
avgpooled_x2 = avgpooled_x2.view(avgpooled_x2.size(0), -1) # bxd
return avgpooled_x, x, avgpooled_x2
|