Working with PyTorch tensors, I need to split the batch of items by its flags, so the items in x_batch_one and x_batch_two are lists of pairs that have the same flags.

def pairs_by_flags(self, x_batch, y_batch, max_flag):
#split x_batch on pairs

    sizes = (x_batch.size(0)//2, *x_batch.size()[1:] )

    items_pos = torch.zeros((max_flag,), dtype = torch.long, device = x_batch.device)

    x_batch_one = torch.zeros(sizes, device = x_batch.device)
    x_batch_one_pos = 0
    x_batch_two = torch.zeros(sizes, device = x_batch.device)

    for index, flag in enumerate(y_batch):
        x_item = x_batch[index]
        i_flag = items_pos[flag]
        if i_flag:                
            x_batch_two[i_flag-1] = x_item
            items_pos[flag] = 0
            x_batch_one[x_batch_one_pos] = x_item
            items_pos[flag] = x_batch_one_pos

    return x_batch_one, x_batch_two    

Here is a simple example with single digits:

from torch_helper.myt import thelp
import torch

arr = torch.tensor([2,7,6,2,7,3,2,8,8,6,2,3])

arr1, arr2 =   thelp.pairs_by_flags(arr, arr, 9)




tensor([2., 7., 6., 3., 2., 8.])
tensor([2., 7., 6., 3., 2., 8.])

But after making this naive solution I wonder if it could be improved, perhaps eliminating the loop using some tensor [:,,,:,,,:] magic.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.