| 62 | |
| 63 | |
| 64 | class AlternateCorrBlock: |
| 65 | def __init__(self, fmap1, fmap2, num_levels=4, radius=4): |
| 66 | self.num_levels = num_levels |
| 67 | self.radius = radius |
| 68 | |
| 69 | self.pyramid = [(fmap1, fmap2)] |
| 70 | for i in range(self.num_levels): |
| 71 | fmap1 = F.avg_pool2d(fmap1, 2, stride=2) |
| 72 | fmap2 = F.avg_pool2d(fmap2, 2, stride=2) |
| 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()) |