How to handle this case in pytorch?
I have the following kinds of two corelated tensors data and mask. The size of data is 1x2x24x2 and the size of mask is 1x2x24. In mask, the True means the corresponding data in data is efficient and should be gradient back-propagated.
data:
tensor([[[[ 1.0663e+03,  5.5993e+02],
          [ 1.0612e+03,  7.2023e+02],
          [ 1.0831e+03,  7.2179e+02],
          [ 1.0945e+03,  5.6083e+02],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10]],

         [[ 6.9314e+02,  1.9700e+02],
          [ 6.3300e+02,  2.6924e+02],
          [ 6.3300e+02,  3.4165e+02],
          [ 7.7515e+02,  4.6000e+02],
          [ 8.2805e+02,  4.6000e+02],
          [ 9.0900e+02,  3.6276e+02],
          [ 9.0900e+02,  2.9035e+02],
          [ 7.9688e+02,  1.9700e+02],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10],
          [-1.0000e+10,  1.0000e+10]]]]) 
torch.Size([1, 2, 24, 2])

mask:
tensor([[[ True,  True,  True,  True, False, False, False, False, False, False,
          False, False, False, False, False, False, False, False, False, False,
          False, False, False, False],
         [ True,  True,  True,  True,  True,  True,  True,  True, False, False,
          False, False, False, False, False, False, False, False, False, False,
          False, False, False, False]]]) 
torch.Size([1, 2, 24])

Based on the data and mask, I need further processing. For example, I need to do this kind of for loop:
B, N = data.shape[0:2]
result = torch.zeros((B, N), dtype=torch.float32)
for b in range(B):
    for n in range(N):
        tmp_data = data[b,n]  # 24 x 2
        tmp_mask = mask[b,n]  # 24
        selected = tmp_data[tmp_mask]
        selected = torch.cat([selected, selected[0][None]], dim=0)
        total = selected[0:-1, 0] * selected[1:, 1] - selected[0:-1, 1] * selected[1:, 0]   # do cross product
        total = torch.sum(total, dim=0)
        result[b,n] = torch.abs(total) / 2

Then, the result tensor is fed into subsequent processing, e.g. loss computation. My question is, is there any speed up method to get rid of the for-loop? The main difficulty lies in that for each b and n, the count of True in mask is not equal, sometimes it is 3, sometimes it is 5, but the maximum count of True in each sample is 8. So, in the for-loop, it is necessary to handle each case one by one. Is there any vectorized way for this?
Another difficulty lies in the gradient back-propagation, which I need further investigation to check if this routine can work correctly for gradient back-propagation.
I'm looking forward to an elegant way to handle this kind of tensor processing in pytorch. Thanks in advance!
UPDATED:
The data and mask were sorted by the previous routine, for example:
sorted_inds = ...  # B x N x 24

data = torch.gather(data_original, dim=2, index=sorted_indices.unsqueeze(-1).repeat(1,1,1,2))  # B x N x 24 x 2
mask= torch.gather(mask_original, dim=2, index=sorted_indices)  # B x N x 24

# then the data and mask are obtained
# continue processing in this post
...


You are looking to vectorize the two for loops with tensor operations. Generally speaking this is possible by just ignoring the first two axes with performing operations on the last two. However, in in this particular case I don't think this is achievable.
The reason lies with the masking that takes place on data with mask (via indexing). This operation will only return a flattened tensor:.
data[mask]
tensor([[1066.3000,  559.9300],
        [1061.2000,  720.2300],
        [1083.1000,  721.7900],
        [1094.5000,  560.8300],
        [ 693.1400,  197.0000],
        [ 633.0000,  269.2400],
        [ 633.0000,  341.6500],
        [ 775.1500,  460.0000],
        [ 828.0500,  460.0000],
        [ 909.0000,  362.7600],
        [ 909.0000,  290.3500],
        [ 796.8800,  197.0000]])

As you can see the correct values are retained but not the desired shape. That's where the problem lies. This is what we would have wished to have (in order to proceed with a vectorized form). This is not possible as all tensor on a dimension must have the same length: the size on that dimension.
tensor([[[1066.3000,  559.9300],
         [1061.2000,  720.2300],
         [1083.1000,  721.7900],
         [1094.5000,  560.8300]],
        [[ 693.1400,  197.0000],
         [ 633.0000,  269.2400],
         [ 633.0000,  341.6500],
         [ 775.1500,  460.0000],
         [ 828.0500,  460.0000],
         [ 909.0000,  362.7600],
         [ 909.0000,  290.3500],
         [ 796.8800,  197.0000]])


Concerning the cross-product. PyTorch provides a function for that: torch.cross, but it can only be used for 3D tensors. Above that it is not possible to use it. I think your implementation is fine.

Some suggestions, these are not groundbreaking improvements but can help:

There's a trick with indexing to avoid having to do x.unsqueeze(0) (or as you did x[None]) on one-element tensors: by slicing x[:1].

Argument dim generally defaults to 0. For .sum() you don't need it since you are looking for the total sum and not a sum per-axis.

You could flatten the first two axes, perform a single loop and reshape at the end.


Here, I have taken the liberty of tweaking your code:
B, N, *_ = data.shape
data_ = data.reshape(B*N, -1, 2)
mask_ = mask.reshape(B*N, -1)

result = []
for i in range(B*N):
    selected = data_[i][mask_[i]]
    looped = torch.cat([selected, selected[:1]])
    cross = looped[:-1, 0]*looped[1:, 1] - looped[:-1, 1]*looped[1:, 0]
    result.append(cross.sum().abs() / 2)

torch.tensor(result).reshape(B, N)

