The Conditional Position Encoding (CPE) module. The CPE is the implementation of 'Conditional Positional Encodings for Vision Transformers '_. Args: in_channels (int): Number of input channels. embed_dims (int): The feature dimension. Def
| 317 | |
| 318 | |
| 319 | class ConditionalPositionEncoding(BaseModule): |
| 320 | """The Conditional Position Encoding (CPE) module. |
| 321 | |
| 322 | The CPE is the implementation of 'Conditional Positional Encodings |
| 323 | for Vision Transformers <https://arxiv.org/abs/2102.10882>'_. |
| 324 | |
| 325 | Args: |
| 326 | in_channels (int): Number of input channels. |
| 327 | embed_dims (int): The feature dimension. Default: 768. |
| 328 | stride (int): Stride of conv layer. Default: 1. |
| 329 | """ |
| 330 | |
| 331 | def __init__(self, in_channels, embed_dims=768, stride=1, init_cfg=None): |
| 332 | super(ConditionalPositionEncoding, self).__init__(init_cfg=init_cfg) |
| 333 | self.proj = nn.Conv2d( |
| 334 | in_channels, embed_dims, kernel_size=3, stride=stride, padding=1, bias=True, groups=embed_dims |
| 335 | ) |
| 336 | self.stride = stride |
| 337 | |
| 338 | def forward(self, x, hw_shape): |
| 339 | b, n, c = x.shape |
| 340 | h, w = hw_shape |
| 341 | feat_token = x |
| 342 | cnn_feat = feat_token.transpose(1, 2).view(b, c, h, w) |
| 343 | if self.stride == 1: |
| 344 | x = self.proj(cnn_feat) + cnn_feat |
| 345 | else: |
| 346 | x = self.proj(cnn_feat) |
| 347 | x = x.flatten(2).transpose(1, 2) |
| 348 | return x |
| 349 | |
| 350 | |
| 351 | @BACKBONES.register_module() |