Visualization of position embeddings about vision_transformer HOT 2 OPEN

def patch_cosSim(model, pos_i, pos_j):
    """
    Arguments:
    model: ViT
    pos_i: ith patch (int)
    pos_j: jth patch (int)

    returns: cos similarity of the pos emb of the (i,j) patch
    """
    cos = nn.CosineSimilarity(dim=0)
    s = model.embeddings.position_embeddings.shape
    pos_patch = model.embeddings.position_embeddings.view(*s[1:])[:-1].view(14,14,-1)
    n_rows = pos_patch.shape[0]
    n_cols = pos_patch.shape[1]
    patches = []
    for i in range(n_rows):
        for j in range(n_cols):
            patches.append(cos(pos_patch[pos_i][pos_j], pos_patch[i][j]).detach().cpu().numpy())
    return patches

i = 0
j = 1
patch = patch_cosSim(i, j)
ax = fig.add_subplot()
ax.set_axis_off()
ax.imshow(np.array(patch).reshape(14,14), cmap='hot', interpolation='nearest')