HaloAttention掩码问题
Author: L-wwwwCreated Dec 19, 2024Updated Dec 19, 2024
mask out padding (in the paper, they claim to not need masks, but what about padding?)
mask = torch.ones(1, 1, h, w, device = device)
mask = F.unfold(mask, kernel_size = block + (halo * 2), stride = block, padding = halo)
mask = repeat(mask, '() j i -> (b i h) () j', b = b, h = heads)
mask = mask.bool()
max_neg_value = -torch.finfo(sim.dtype).max
sim.masked_fill_(mask, max_neg_value)
代码似乎将所有非0值置为无穷小,根据注释似乎是想将padding位置 置为无穷小
Source: xmu-xiaoma666/External-Attention-pytorch