Spaces:
Runtime error
Runtime error
| # 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 | |
| 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 | |