Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,526 Bytes
a249588 |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
# Copyright (c) OpenMMLab. All rights reserved.
import math
import torch
import torch.nn.functional as F
from mmcv.cnn import Linear
from mmengine.model import BaseModule, ModuleList
from torch import Tensor
class FFN(BaseModule):
"""Very simple multi-layer perceptron with relu. Mostly used in DETR series
detectors.
Args:
input_dim (int): Feature dim of the input tensor.
hidden_dim (int): Feature dim of the hidden layer.
output_dim (int): Feature dim of the output tensor.
num_layers (int): Number of FFN layers..
"""
def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
num_layers: int) -> None:
super().__init__()
self.num_layers = num_layers
self.layers = ModuleList()
self.layers.append(Linear(input_dim, hidden_dim))
for _ in range(num_layers - 2):
self.layers.append(Linear(hidden_dim, hidden_dim))
self.layers.append(Linear(hidden_dim, output_dim))
def forward(self, x: Tensor) -> Tensor:
"""Forward function of FFN.
Args:
x (Tensor): The input feature, has shape
(num_queries, bs, input_dim).
Returns:
Tensor: The output feature, has shape
(num_queries, bs, output_dim).
"""
for i, layer in enumerate(self.layers):
x = layer(x)
if i < self.num_layers - 1:
x = F.relu(x)
return x
class PositionEmbeddingSineHW(BaseModule):
"""This is a more standard version of the position embedding, very similar
to the one used by the Attention is all you need paper, generalized to work
on images."""
def __init__(self,
num_pos_feats=64,
temperatureH=10000,
temperatureW=10000,
normalize=False,
scale=None):
super().__init__()
self.num_pos_feats = num_pos_feats
self.temperatureH = temperatureH
self.temperatureW = temperatureW
self.normalize = normalize
if scale is not None and normalize is False:
raise ValueError('normalize should be True if scale is passed')
if scale is None:
scale = 2 * math.pi
self.scale = scale
def forward(self, mask: Tensor):
assert mask is not None
not_mask = ~mask
y_embed = not_mask.cumsum(1, dtype=torch.float32)
x_embed = not_mask.cumsum(2, dtype=torch.float32)
if self.normalize:
eps = 1e-6
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
dim_tx = torch.arange(
self.num_pos_feats, dtype=torch.float32, device=mask.device)
dim_tx = self.temperatureW**(2 * (dim_tx // 2) / self.num_pos_feats)
pos_x = x_embed[:, :, :, None] / dim_tx
dim_ty = torch.arange(
self.num_pos_feats, dtype=torch.float32, device=mask.device)
dim_ty = self.temperatureH**(2 * (dim_ty // 2) / self.num_pos_feats)
pos_y = y_embed[:, :, :, None] / dim_ty
pos_x = torch.stack(
(pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()),
dim=4).flatten(3)
pos_y = torch.stack(
(pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()),
dim=4).flatten(3)
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
return pos
|