Split tensor based on segments

vision
1

Hi!

I am looking for the most efficient way to split a tensor into multiple different tensors based on a 1-D index tensor.

Here is the full problem description:

I have a point cloud tensor with dimension (batchsize, 3, 2048) and an index tensor of size (batchsize, 2048) (LongTensor). This index array contains integer values from 0 to 3 (inclusive). I want to split the point cloud based on those indices such that I end up with 4 different tensors of shape (batchsize, 3, N_i) where the sum over all 4 different N_i equals 2048.

I’ve looked into torch.index_select but I can’t figure out, how I could use this function to my advantage.

Your help would be much appreciated!

2

I’m not sure to understand the use case completely.
Could you post some examples using smaller tensor shapes, e.g. [1, 3, 10], and demonstrate how you would like to split this tensor using the indices?

3

Coming here having a similar question. If I understood the OP correctly, given a tensor A of
(1, 3, 5) and a corresponding tensor B of shape (5,) where for example

A=
[[[3, 0, 0, 1, 8, 9, 5, 4, 1, 6],
[0, 0, 0, 8, 0, 7, 1, 0, 0, 8],
[1, 0, 2, 5, 3, 6, 3, 2, 4, 7]]]

and B=[0, 0, 1, 1, 0, 2, 2, 3, 3, 3]

the goal is to get 4 tensors where the first contains all elements of A whose corresponding B is 0.
Thus
[[[3, 0, 8],
[0, 0, 0],
[1, 0, 3]]]

the second one would be
[[[0, 1],
[0, 8],
[2, 5]]]

and so forth.

Is there an easy was of doing it. And to answer my actual question, is there also a way to generalize it to arbitrary index values and then split in a similar fashion such that each resulting tensor contains values with the same corresponding index?