Hi All,
I have two 3D tensors X and Q of shape (5, 16, 128) on which I do cosine similarity on 2nd dim to get a (5, 16) cosine-similarity vector. I then sort this cosine-similarity vector, to get indices of most-to-least similar vectors in the original vector Q. In code
X = torch.ones((5, 16, 128))
Q = torch.randn(5, 16, 128) # dim=2 is features
neg_samples = 11
pos_samples = 5
cos_sim = torch.nn.CosineSimilarity(dim=2, eps=1e-6)
sim_scores = cos_sim(X, Q) # sim_scores output shape: [5,16]
_, sorted_idxs = torch.topk(sim_scores, k=sim_scores.shape[1], sorted=True, dim=1) #sort sim_scores to get sorted indexes same shape [5, 16]
#since sorted_idxs is sorted descending on similarity, we can slice to divide pos/neg samples.
pos_idxs = sorted_idxs[:, :pos_samples] #gives [5, 5] most positive idxs from sorted_idxs
neg_idxs = sorted_idxs[:, pos_samples:(pos_samples+neg_samples)] #gives [5, 11] most negative idxs from sorted_idxs
How do i retrieve the corresponding tensors from Q which is of shape [5, 16, 128] using these pos_idxs shape [5,5] and neg_idxs shape [5,11]?
I tried Q[ pos_idxs, :] but i get output tensor of shape [5, 5, 16, 128]. I want to slice Q into [5, 5, 128] and [5, 11, 128] two separate tensors based on 2D pos_idxs and neg_idxs index tensors.
How do i do this in pytorch using tensor operations?
thanks