| 73 | self.pyramid.append((fmap1, fmap2)) |
| 74 | |
| 75 | def __call__(self, coords): |
| 76 | coords = coords.permute(0, 2, 3, 1) |
| 77 | B, H, W, _ = coords.shape |
| 78 | dim = self.pyramid[0][0].shape[1] |
| 79 | |
| 80 | corr_list = [] |
| 81 | for i in range(self.num_levels): |
| 82 | r = self.radius |
| 83 | fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous() |
| 84 | fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous() |
| 85 | |
| 86 | coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous() |
| 87 | corr, = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r) |
| 88 | corr_list.append(corr.squeeze(1)) |
| 89 | |
| 90 | corr = torch.stack(corr_list, dim=1) |
| 91 | corr = corr.reshape(B, -1, H, W) |
| 92 | return corr / torch.sqrt(torch.tensor(dim).float()) |