Positional encoding using random spatial frequencies.
| 132 | |
| 133 | |
| 134 | class PositionEmbeddingRandom(nn.Module): |
| 135 | """ |
| 136 | Positional encoding using random spatial frequencies. |
| 137 | """ |
| 138 | |
| 139 | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: |
| 140 | super().__init__() |
| 141 | if scale is None or scale <= 0.0: |
| 142 | scale = 1.0 |
| 143 | self.register_buffer( |
| 144 | "positional_encoding_gaussian_matrix", |
| 145 | scale * torch.randn((2, num_pos_feats)), |
| 146 | ) |
| 147 | |
| 148 | @torch.no_grad() |
| 149 | def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: |
| 150 | """Positionally encode points that are normalized to [0,1].""" |
| 151 | # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape |
| 152 | coords = 2 * coords - 1 |
| 153 | coords = coords @ self.positional_encoding_gaussian_matrix.to(coords.dtype) |
| 154 | coords = 2 * np.pi * coords |
| 155 | # outputs d_1 x ... x d_n x C shape |
| 156 | return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) |
| 157 | |
| 158 | @torch.no_grad() |
| 159 | def forward(self, size: Tuple[int, int]) -> torch.Tensor: |
| 160 | """Generate positional encoding for a grid of the specified size.""" |
| 161 | h, w = size |
| 162 | device = self.positional_encoding_gaussian_matrix.device |
| 163 | |
| 164 | # Force fp32 (https://github.com/huggingface/transformers/pull/29285) |
| 165 | with torch.autocast(device_type=device.type, enabled=False): |
| 166 | grid = torch.ones((h, w), device=device, dtype=torch.float32) |
| 167 | y_embed = grid.cumsum(dim=0) - 0.5 |
| 168 | x_embed = grid.cumsum(dim=1) - 0.5 |
| 169 | y_embed = y_embed / h |
| 170 | x_embed = x_embed / w |
| 171 | pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) |
| 172 | |
| 173 | pe = pe.to(self.positional_encoding_gaussian_matrix.dtype) |
| 174 | return pe.permute(2, 0, 1) # C x H x W |
| 175 | |
| 176 | @torch.no_grad() |
| 177 | def forward_with_coords(self, coords_input: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor: |
| 178 | """Positionally encode points that are not normalized to [0,1].""" |
| 179 | assert coords_input.dtype == torch.float, 'coords_input must be in float32' |
| 180 | |
| 181 | # Force fp32 (https://github.com/huggingface/transformers/pull/29285) |
| 182 | with torch.autocast(device_type=coords_input.device.type, enabled=False): |
| 183 | coords = coords_input.clone() |
| 184 | coords[:, :, 0] = coords[:, :, 0] / image_size[1] |
| 185 | coords[:, :, 1] = coords[:, :, 1] / image_size[0] |
| 186 | pe = self._pe_encoding(coords.to(torch.float)) # B x N x C |
| 187 | |
| 188 | pe = pe.to(self.positional_encoding_gaussian_matrix.dtype) |
| 189 | return pe |
| 190 | |
| 191 |