Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import math | |
| ###################################################################################### | |
| # position embedding | |
| ###################################################################################### | |
| class PositionEmbeddingLearned(nn.Module): | |
| """ | |
| This is a learned version of the position embedding | |
| """ | |
| def __init__(self, num_pos_feats=256): | |
| super().__init__() | |
| self.row_embed = nn.Embedding(32, num_pos_feats) | |
| self.col_embed = nn.Embedding(32, num_pos_feats) | |
| self.reset_parameters() | |
| def reset_parameters(self): | |
| nn.init.uniform_(self.row_embed.weight) | |
| nn.init.uniform_(self.col_embed.weight) | |
| def forward(self, x, mask): | |
| h, w = x.shape[-2:] | |
| i = torch.arange(w, device=x.device) | |
| j = torch.arange(h, device=x.device) | |
| x_emb = self.col_embed(i).unsqueeze(0).repeat(h, 1, 1) | |
| y_emb = self.row_embed(j).unsqueeze(1).repeat(1, w, 1) | |
| pos = (x_emb + y_emb).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1) | |
| return pos | |
| class PositionEmbeddingSine(nn.Module): | |
| """ | |
| This is a standard version of the position embedding, very similar to the one used by the | |
| "Attention is all you need" paper, generalized to work on examples | |
| """ | |
| def __init__(self, feats_dim=512, temperature=10000, normalize=False, scale=None): | |
| """ | |
| explicitly encode the position using the sinusoid: | |
| PE(pos,2i) = sin(pos/temperature^(2*i/d_model)) | |
| PE(pos,2i+1) = cos(pos/temperature^(2*i/d_model)) | |
| :param feats_dim: the dimension of features, each dimension of the positional embedding to a sinusoid | |
| :param temperature: wavelengths from a geometric progression from scale | |
| :param normalize: whether to normalize the position to (0,1) | |
| :param scale: scale for the position embedding | |
| """ | |
| super(PositionEmbeddingSine, self).__init__() | |
| self.feats_dim = feats_dim | |
| self.T = temperature | |
| self.norm = normalize | |
| if scale is None: | |
| scale = 2 * math.pi | |
| self.scale = scale | |
| def forward(self, x, mask): | |
| x_embed = mask.cumsum(1, dtype=torch.float32) | |
| y_embed = mask.cumsum(2, dtype=torch.float32) | |
| if self.norm: | |
| eps = 1e-5 | |
| x_embed = x_embed / (x_embed[:, -1:, :] + eps) * self.scale | |
| y_embed = y_embed / (y_embed[:, :, -1:] + eps) * self.scale | |
| dim_t = torch.arange(self.feats_dim, dtype=torch.float32, device=x.device) | |
| dim_t = self.T ** (2*(dim_t//2)/self.feats_dim) | |
| pos_x = x_embed[:, :, :, None] / dim_t | |
| pos_y = y_embed[:, :, :, None] / dim_t | |
| pos_x[:, :, :, 0::2], pos_x[:, :, :, 1::2] = pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos() | |
| pos_y[:, :, :, 0::2], pos_y[:, :, :, 1::2] = pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos() | |
| pos = (pos_x + pos_y).permute(0, 3, 1, 2) * 0.5 | |
| return pos | |
| def build_position_embed(embed_type='learned', feats_dim=512, temperature=10000): | |
| if embed_type == 'sine': | |
| pos_embed = PositionEmbeddingSine(feats_dim, temperature, normalize=True) | |
| elif embed_type == 'learned': | |
| pos_embed = PositionEmbeddingLearned(feats_dim) | |
| else: | |
| raise ValueError(f"nor supported {embed_type}") | |
| return pos_embed | |