| import torch |
| from torch import nn |
|
|
|
|
| class PositionalEncodingsFixed(nn.Module): |
|
|
| def __init__(self, emb_dim, temperature=10000): |
|
|
| super(PositionalEncodingsFixed, self).__init__() |
|
|
| self.emb_dim = emb_dim |
| self.temperature = temperature |
|
|
| def _1d_pos_enc(self, mask, dim): |
| temp = torch.arange(self.emb_dim // 2).float().to(mask.device) |
| temp = self.temperature ** (2 * (temp.div(2, rounding_mode='floor')) / self.emb_dim) |
|
|
| enc = (~mask).cumsum(dim).float().unsqueeze(-1) / temp |
| enc = torch.stack([ |
| enc[..., 0::2].sin(), enc[..., 1::2].cos() |
| ], dim=-1).flatten(-2) |
|
|
| return enc |
|
|
| def forward(self, bs, h, w, device): |
| mask = torch.zeros(bs, h, w, dtype=torch.bool, requires_grad=False, device=device) |
| x = self._1d_pos_enc(mask, dim=2) |
| y = self._1d_pos_enc(mask, dim=1) |
|
|
| return torch.cat([y, x], dim=3).permute(0, 3, 1, 2) |
|
|
