File size: 3,112 Bytes
f1c1609 |
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 |
# ------------------------------------------------------------------------
# Deformable DETR
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------
"""
Various positional encodings for the transformer.
"""
import math
import torch
from torch import nn
from misc.detr_utils.misc import NestedTensor
class PositionEmbeddingSine(nn.Module):
"""
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, temperature=10000, normalize=False, scale=None):
super().__init__()
self.num_pos_feats = num_pos_feats
self.temperature = temperature
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
self.max_duration = 256
self.duration_embed_layer = nn.Linear(self.max_duration, self.max_duration)
def forward(self, tensor_list: NestedTensor):
x = tensor_list.tensors
mask = tensor_list.mask
duration = tensor_list.duration
assert mask is not None
not_mask = ~mask
x_embed = not_mask.cumsum(1, dtype=torch.float32)
if self.normalize:
eps = 1e-6
x_embed = (x_embed - 0.5) / (x_embed[:, -1:] + eps) * self.scale
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
# dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
dim_t = self.temperature ** (2 * (torch.div(dim_t, 2, rounding_mode='floor')) / self.num_pos_feats)
pos_x = x_embed[:, :, None] / dim_t
pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
dur_embed = self.duration_embedding(duration).reshape(-1,1,self.max_duration).expand_as(pos_x)
pos = torch.cat((pos_x, dur_embed), dim=2).permute(0, 2, 1)
return pos
def duration_embedding(self, durations):
out = torch.zeros(len(durations), self.max_duration, device=durations.device)
durations = durations.int()
for ii in range(len(durations)):
out[ii, :durations[ii]] = 1
out = self.duration_embed_layer(out)
return out
def build_position_encoding(position_embedding, N_steps):
if position_embedding in ('v2', 'sine'):
# TODO find a better way of exposing other arguments
position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
else:
raise ValueError(f"not supported {position_embedding}")
return position_embedding
|