Shuffling two 2D tensors in PyTorch and maintaining same order correlation

Question:

Is it possible to shuffle two 2D tensors in PyTorch by their rows, but maintain the same order for both? I know you can shuffle a 2D tensor by rows with the following code:

a=a[torch.randperm(a.size()[0])]

To elaborate:
If I had 2 tensors

a = torch.tensor([[1, 1, 1, 1, 1],
            [2, 2, 2, 2, 2],
            [3, 3, 3, 3, 3]])

b = torch.tensor([[4, 4, 4, 4, 4],
            [5, 5, 5, 5, 5],
            [6, 6, 6, 6, 6]])

And ran them through some function/block of code to shuffle randomly but maintain correlation and produce something like the following

a = torch.tensor([[2, 2, 2, 2, 2],
            [1, 1, 1, 1, 1],
            [3, 3, 3, 3, 3]])

b = torch.tensor([[5, 5, 5, 5, 5],
            [4, 4, 4, 4, 4],
            [6, 6, 6, 6, 6]])

My current solution is converting to a list, using the random.shuffle() function like below.

a_list = a.tolist()
b_list = b.tolist()
temp_list = list(zip(a_list , b_list ))
random.shuffle(temp_list) # Shuffle
a_temp, b_temp = zip(*temp_list)
a_list, b_list = list(a_temp), list(b_temp)
            
# Convert back to tensors
a = torch.tensor(a_list)
b = torch.tensor(b_list)

This takes quite a while and was wondering if there is a better way.

Asked By: zahlawi

||

Answers:

You mean

indices = torch.randperm(a.size()[0])
a=a[indices]
b=b[indices]

?

Answered By: hkchengrex
Categories: questions Tags: , , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.