Pytorch cut 2d array by lengths on first dimension
I have a 2d array, let's say of size torch.tensor(batch_size, 1000).  The 1000 array from the second dimension is actually variable length.  I have a second array of size [batch_size] containing the length for each rows...

Here is an example code snippet:

# preds is the 2d array of size [batch_size, 1000]
# lengths is a 1d array containing the lengths of each row of preds
res_pred = []
for i in range(len(preds)):
    length = lengths[i].item()
    res_pred += [preds[i][:length]]

result = torch.cat(res_pred).flatten()


I do the same thing for my targets and then I can apply a loss function to both.

I was wondering if there was a single vectorized operation I could do to extract all batch_size vectors of variable lengths and torch.cat them together. Right now I am looping on the first dimension, but this feel slow.

Thanks,

You can create a 2D mask tensor with the number of True's in i-th row given by lengths[i]. Here's one example:

batch_size = 6
n = 5

preds = torch.arange(batch_size * n).reshape(batch_size, n)
# tensor([[ 0,  1,  2,  3,  4],
#         [ 5,  6,  7,  8,  9],
#         [10, 11, 12, 13, 14],
#         [15, 16, 17, 18, 19],
#         [20, 21, 22, 23, 24],
#         [25, 26, 27, 28, 29]])

#lengths = np.random.randint(0, n+1, batch_size)
lengths = torch.randint(0, n+1, (batch_size, ))
# tensor([2, 0, 5, 3, 3, 2])


Let's create the mask and get our result (probably there is a better way to create such a mask, but that's what I came up with):

#mask = np.tile(range(n), (batch_size,1)) &lt; lengths[:,None]
mask = torch.arange(n).repeat((batch_size,1)) &lt; lengths[:, None]
# tensor([[ True,  True, False, False, False],
#        [False, False, False, False, False],
#        [ True,  True,  True,  True,  True],
#        [ True,  True,  True, False, False],
#        [ True,  True,  True, False, False],
#        [ True,  True, False, False, False]])

#result = preds[mask]
result = torch.masked_select(preds, mask)
# tensor([0, 1, 10, 11, 12, 13, 14, 15, 16, 17, 20, 21, 22, 25, 26])


This produces the same result as your code:

res_pred = []
for i in range(len(preds)):
    length = lengths[i].item()
    res_pred += [preds[i][:length]]

result = torch.cat(res_pred).flatten()

