Shuffle Two PyTorch Tensors the Same Way

At times in Pytorch it might be useful to shuffle two separate tensors in the same way, with the result that the shuffled elements create two new tensors which maintain the pairing of elements between the tensors. An example might be to shuffle a dataset and ensure the labels are still matched correctly after the shuffling.

Import Torch

We only need torch for this, it is possible to achieve this is a very similar way in numpy, but I prefer to use Pytorch for simplicity.

import torch

Import Create Example Tensors

Here we create two tensors one runs 0 to 9 and the second runs 9 to 0.

x = torch.arange(10)
y = torch.arange(9, -1, -1)
print(x)
print(y)

tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
tensor([9, 8, 7, 6, 5, 4, 3, 2, 1, 0])

These new tensor elements are tensors, and are paired as follows, the next steps will shuffle the position of these elements while maintaining their pairing.

0->9, 1->8, 2->7, 3->6, 4->5, 5->4, 6->3, 7->2, 8->1, 9->0

Perform Shuffle

Using the index and original tensors we now create two new shuffled tensors.

If using numpy we can achieve the same thing using np.arange(x.shape[0]) then np.random.shuffle(shuffle_index)

shuffle_index = torch.randperm(x.shape[0])
shuffle_index

tensor([9, 3, 2, 8, 4, 5, 0, 7, 1, 6])

Perform Shuffle

Using the index and original tensors we now create two new shuffled tensors.

x_new = x[shuffle_index]
y_new = y[shuffle_index]
print(x_new)
print(y_new)

tensor([9, 3, 2, 8, 4, 5, 0, 7, 1, 6])
tensor([0, 6, 7, 1, 5, 4, 9, 2, 8, 3])

Final Thoughts

The new output tensors are now shuffled but shuffled in the same way maintaining the relationships of elements between them.

Care should be taken if using multi-dimensional arrays to only shuffle the index not all the data (unless that is your intent).