ZF/vit-pytorch
1
0
Fork 0
mirror of https://github.com/lucidrains/vit-pytorch.git synced 2025-12-30 08:02:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

maxvit intent to build (#211)

Browse Source 
complete hybrid mbconv + block / grid efficient self attention MaxViT
This commit is contained in:
Phil Wang
2022-04-06 16:12:17 -07:00
committed by GitHub
parent 946b19be64
commit c7bb5fc43f
7 changed files with 317 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vit_pytorch/levit.py
View File

@@ -71,8 +71,8 @@ class Attention(nn.Module):
q_range = torch.arange(0, fmap_size, step = (2 if downsample else 1))
k_range = torch.arange(fmap_size)
q_pos = torch.stack(torch.meshgrid(q_range, q_range), dim = -1)
k_pos = torch.stack(torch.meshgrid(k_range, k_range), dim = -1)
q_pos = torch.stack(torch.meshgrid(q_range, q_range, indexing = 'ij'), dim = -1)
k_pos = torch.stack(torch.meshgrid(k_range, k_range, indexing = 'ij'), dim = -1)
q_pos, k_pos = map(lambda t: rearrange(t, 'i j c -> (i j) c'), (q_pos, k_pos))
rel_pos = (q_pos[:, None, ...] - k_pos[None, :, ...]).abs()