(self, x)
| 246 | return ref |
| 247 | |
| 248 | def forward(self, x): |
| 249 | |
| 250 | B, C, H, W = x.size() |
| 251 | dtype, device = x.dtype, x.device |
| 252 | |
| 253 | q = self.proj_q(x) |
| 254 | q_off = einops.rearrange(q, 'b (g c) h w -> (b g) c h w', g=self.n_groups, c=self.n_group_channels) |
| 255 | offset = self.conv_offset(q_off).contiguous() # B * g 2 Hg Wg |
| 256 | Hk, Wk = offset.size(2), offset.size(3) |
| 257 | n_sample = Hk * Wk |
| 258 | |
| 259 | if self.offset_range_factor >= 0 and not self.no_off: |
| 260 | offset_range = torch.tensor([1.0 / (Hk - 1.0), 1.0 / (Wk - 1.0)], device=device).reshape(1, 2, 1, 1) |
| 261 | offset = offset.tanh().mul(offset_range).mul(self.offset_range_factor) |
| 262 | |
| 263 | offset = einops.rearrange(offset, 'b p h w -> b h w p') |
| 264 | reference = self._get_ref_points(Hk, Wk, B, dtype, device) |
| 265 | |
| 266 | if self.no_off: |
| 267 | offset = offset.fill_(0.0) |
| 268 | |
| 269 | if self.offset_range_factor >= 0: |
| 270 | pos = offset + reference |
| 271 | else: |
| 272 | pos = (offset + reference).clamp(-1., +1.) |
| 273 | |
| 274 | if self.no_off: |
| 275 | x_sampled = F.avg_pool2d(x, kernel_size=self.stride, stride=self.stride) |
| 276 | assert x_sampled.size(2) == Hk and x_sampled.size(3) == Wk, f"Size is {x_sampled.size()}" |
| 277 | else: |
| 278 | x_sampled = F.grid_sample( |
| 279 | input=x.reshape(B * self.n_groups, self.n_group_channels, H, W), |
| 280 | grid=pos[..., (1, 0)], # y, x -> x, y |
| 281 | mode='bilinear', align_corners=True) # B * g, Cg, Hg, Wg |
| 282 | |
| 283 | |
| 284 | x_sampled = x_sampled.reshape(B, C, 1, n_sample) |
| 285 | |
| 286 | q = q.reshape(B * self.n_heads, self.n_head_channels, H * W) |
| 287 | k = self.proj_k(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample) |
| 288 | v = self.proj_v(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample) |
| 289 | |
| 290 | attn = torch.einsum('b c m, b c n -> b m n', q, k) # B * h, HW, Ns |
| 291 | attn = attn.mul(self.scale) |
| 292 | |
| 293 | if self.use_pe and (not self.no_off): |
| 294 | |
| 295 | if self.dwc_pe: |
| 296 | residual_lepe = self.rpe_table(q.reshape(B, C, H, W)).reshape(B * self.n_heads, self.n_head_channels, H * W) |
| 297 | elif self.fixed_pe: |
| 298 | rpe_table = self.rpe_table |
| 299 | attn_bias = rpe_table[None, ...].expand(B, -1, -1, -1) |
| 300 | attn = attn + attn_bias.reshape(B * self.n_heads, H * W, n_sample) |
| 301 | elif self.log_cpb: |
| 302 | q_grid = self._get_q_grid(H, W, B, dtype, device) |
| 303 | displacement = (q_grid.reshape(B * self.n_groups, H * W, 2).unsqueeze(2) - pos.reshape(B * self.n_groups, n_sample, 2).unsqueeze(1)).mul(4.0) # d_y, d_x [-8, +8] |
| 304 | displacement = torch.sign(displacement) * torch.log2(torch.abs(displacement) + 1.0) / np.log2(8.0) |
| 305 | attn_bias = self.rpe_table(displacement) # B * g, H * W, n_sample, h_g |
nothing calls this directly
no test coverage detected