Hi, I need help in vectorizing this code. Here, E is (B, 8, 256, 256) vector. K is 7 (knn’s K nearest neighbor). Basically, I want to slide a window on E, and for each pixel find the K nearest pixel vectors in the corresponding window. Right now I have two loops for doing the same procedure, but it takes a very long time.
Z = torch.zeros_like(E)
Z = repeat(Z, 'b c h w -> b (repeat c) h w', repeat=self.K)
### Vectorize this code
for i in range(H):
for j in range(W):
query = E[:, :, i, j] # B, C, 1, 1
tl_x = max(0, i - self.patch_size)
tl_y = max(0, j - self.patch_size)
br_x = min(H, i + self.patch_size)
br_y = min(W, j + self.patch_size)
neighbor_patch = E[:, :, tl_x:br_x+1, tl_y:br_y+1] # B, C, Ph, Pw
query = rearrange(query, 'b c -> b 1 c')
neighbor_patch = rearrange(neighbor_patch, 'b c h w -> b (h w) c')
distance = -torch.cdist(query, neighbor_patch) # b 1 hw
distance_softmax = F.softmax(distance, dim=-1)
_, indices = torch.topk(distance_softmax, dim=-1, k=self.K + 1) # b 1 k
indices = repeat(indices.squeeze(1), 'b k -> b k c', c=C)
neighbors = neighbor_patch.gather(dim=1, index=indices) # b k c
neighbors = neighbors[:, :, 1:].reshape(B, -1)
Z[:, :, i, j] = neighbors
I tried using F.unfold for doing the vectorizing, but I am having difficulty implementing it.