Source code for mmpretrain.models.necks.densecl_neck

# 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

[docs]@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 <>`_. 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))
[docs] 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
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.