Positional encoding using random spatial frequencies.
| 113 | |
| 114 | |
| 115 | class PositionEmbeddingRandom(nn.Module): |
| 116 | """ |
| 117 | Positional encoding using random spatial frequencies. |
| 118 | """ |
| 119 | |
| 120 | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: |
| 121 | super().__init__() |
| 122 | if scale is None or scale <= 0.0: |
| 123 | scale = 1.0 |
| 124 | self.register_buffer( |
| 125 | "positional_encoding_gaussian_matrix", |
| 126 | scale * torch.randn((2, num_pos_feats)), |
| 127 | ) |
| 128 | |
| 129 | def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: |
| 130 | """Positionally encode points that are normalized to [0,1].""" |
| 131 | # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape |
| 132 | coords = 2 * coords - 1 |
| 133 | coords = coords @ self.positional_encoding_gaussian_matrix |
| 134 | coords = 2 * np.pi * coords |
| 135 | # outputs d_1 x ... x d_n x C shape |
| 136 | return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) |
| 137 | |
| 138 | def forward(self, size: Tuple[int, int]) -> torch.Tensor: |
| 139 | """Generate positional encoding for a grid of the specified size.""" |
| 140 | h, w = size |
| 141 | device: Any = self.positional_encoding_gaussian_matrix.device |
| 142 | grid = torch.ones((h, w), device=device, dtype=torch.float32) |
| 143 | y_embed = grid.cumsum(dim=0) - 0.5 |
| 144 | x_embed = grid.cumsum(dim=1) - 0.5 |
| 145 | y_embed = y_embed / h |
| 146 | x_embed = x_embed / w |
| 147 | |
| 148 | pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) |
| 149 | return pe.permute(2, 0, 1) # C x H x W |
| 150 | |
| 151 | def forward_with_coords( |
| 152 | self, coords_input: torch.Tensor, image_size: Tuple[int, int] |
| 153 | ) -> torch.Tensor: |
| 154 | """Positionally encode points that are not normalized to [0,1].""" |
| 155 | coords = coords_input.clone() |
| 156 | coords[:, :, 0] = coords[:, :, 0] / image_size[1] |
| 157 | coords[:, :, 1] = coords[:, :, 1] / image_size[0] |
| 158 | return self._pe_encoding(coords.to(torch.float)) # B x N x C |
| 159 | |
| 160 | |
| 161 | # Rotary Positional Encoding, adapted from: |