Sorted a list of list of tensor based on length

user_2130515 0
Viewed 24
import torch

source = [torch.tensor([1, 2, 3]), torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3, 5, 6])]

source_len = [torch.tensor([3]), torch.tensor([4]), torch.tensor([5])]

source_txt = ["A B C", "A B C D", "A B C D E"]


target =  [torch.tensor([1, 2, 3, 1]), torch.tensor([1, 2, 3]), torch.tensor([1, 2, 3, 5, 6, 2])]

target_len = [torch.tensor([4]), torch.tensor([3]), torch.tensor([6])]

target_txt = ["E F G H", "E F G", "E F G H I J"]

Here each tensor in source corresponds to each tensor in target.

I want to implement kind of "sort within batch" in reverse according to the length of source:

# result should be:

source = [torch.tensor([1, 2, 3, 5, 6]), torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3])]

source_len = [torch.tensor([5]), torch.tensor([4]), torch.tensor([3])]

source_txt = ["A B C D E", "A B C D", "A B C"]


target =  [torch.tensor([1, 2, 3, 5, 6, 2]), torch.tensor([1, 2, 3])torch.tensor([1, 2, 3, 1])]

target_len = [torch.tensor([6]), torch.tensor([3]), torch.tensor([4])]

target_txt = ["E F G H I J", "E F G", "E F G H"]
pytorch list numpy python
1 Answers

You can simply use the second argument returned from torch.sort(), as such:

source_len_sorted ,idx = torch.sort(torch.as_tensor(source_len),descending=True)
source_len_sorted = list(source_len_sorted)
source_sorted = [source[i] for i in idx]
source_txt_sorted = [source_txt[i] for i in idx]

target_sorted =  [target[i] for i in idx]
target_len_sorted = [target_len[i] for i in idx]
target_txt_sorted = [target_txt[i] for i in idx]

If the values were stored in torch tensors rather than lists, you could just index them by idx:

source_sorted = source[idx]

etc etc.

user_104757620
Related