(fmap1, fmap2)
| 51 | |
| 52 | @staticmethod |
| 53 | def corr(fmap1, fmap2): |
| 54 | batch, dim, ht, wd = fmap1.shape |
| 55 | batch, dim, ht2, wd2 = fmap2.shape |
| 56 | fmap1 = fmap1.view(batch, dim, ht*wd) |
| 57 | fmap2 = fmap2.view(batch, dim, ht2*wd2) |
| 58 | |
| 59 | corr = torch.matmul(fmap1.transpose(1,2), fmap2) |
| 60 | corr = corr.view(batch, ht, wd, 1, ht2, wd2) |
| 61 | return corr / torch.sqrt(torch.tensor(dim).float()) |
| 62 | |
| 63 | |
| 64 | class AlternateCorrBlock: |