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