How to multiply a set of masks over an array of n matrices or tensors in python without using loops?

user_11247393 0
Viewed 296

I have N binary masks, along with N x M matrices. I wish to apply the ith mask onto the M matrices at the ith index of the matrix array. Their data types can be either torch tensors or numpy arrays. To illustrate,

If I have an array arr:

arr = torch.rand((2, 3, 3, 3))

or

arr = tensor([[[[0.2336, 0.4841, 0.4121],
          [0.9342, 0.8496, 0.8332],
          [0.4670, 0.8158, 0.7891]],

         [[0.5791, 0.2391, 0.8501],
          [0.9811, 0.0087, 0.0655],
          [0.6587, 0.3105, 0.0931]],

         [[0.8892, 0.8104, 0.9181],
          [0.1605, 0.5280, 0.0905],
          [0.2149, 0.8851, 0.7125]]],


        [[[0.9969, 0.8589, 0.7479],
          [0.4013, 0.5922, 0.0252],
          [0.9267, 0.8123, 0.0711]],

         [[0.7931, 0.6477, 0.0947],
          [0.5969, 0.7751, 0.5662],
          [0.1785, 0.0310, 0.9135]],

         [[0.1490, 0.3623, 0.3670],
          [0.3710, 0.7887, 0.1310],
          [0.2052, 0.0244, 0.6891]]]])

and I generate a mask using:

mask = arr[:, 0, 0:, :] > 0.5

For example:

mask = tensor([[[False, False, False],
         [ True,  True,  True],
         [False,  True,  True]],

        [[ True,  True,  True],
         [False,  True, False],
         [ True,  True, False]]])

i.e. of shape (2, 3, 3). For each set of 3 x 3 x 3 matrices, I want to multiply the corresponding single mask, i.e.

result = mask[i] * arr[i, :, :, :]

or result should ideally be:

result = tensor([[[[0, 0, 0],
          [0.9342, 0.8496, 0.8332],
          [0, 0.8158, 0.7891]],

         [[0, 0, 0],
          [0.9811, 0.0087, 0.0655],
          [0, 0.3105, 0.0931]],

         [[0, 0, 0],
          [0.1605, 0.5280, 0.0905],
          [0, 0.8851, 0.7125]]],


        [[[0.9969, 0.8589, 0.7479],
          [0, 0.5922, 0],
          [0.9267, 0.8123, 0]],

         [[0.7931, 0.6477, 0.0947],
          [0, 0.7751, 0],
          [0.1785, 0.0310, 0]],

         [[0.1490, 0.3623, 0.3670],
          [0, 0.7887, 0],
          [0.2052, 0.0244, 0]]]])

I am using this in a relatively long running piece of code, hence I want to avoid using loops which might blow up the code. I have tried using np.multiply, np.dot and np.matmul, but I keep facing dimensionality issues. Can anyone help me with this?

The following code shows the operation I want to perform, removing the for loop:

for i, m in enumerate(arr):
  result[i] = mask[i] * arr[i]
pytorch matrix numpy arrays python
1 Answers

By using an integer indexing, you removed one of the dimension from arr (2nd axis), and you can't broadcast an array with shape (2, 3, 3) with (2, 3, 3, 3). To make them compatible again, you can add the dimension back by reshaping:

mask.reshape((2, 1, 3, 3)) * arr

Or keep the dimension from the first place by using slice:

mask = arr[:, :1, :, :] > 0.5

and then you can just do mask * arr.

user_49834500
Related