How to fill 2D binary tensor without using for loop?

vision
1

Suppose I have a tensor a and I want to fill the inner with all 1 like tensor b

print(a)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 1., 0., 1., 1., 1., 1., 1., 0., 0.],
        [0., 1., 0., 0., 1., 0., 0., 0., 1., 0.],
        [0., 1., 0., 1., 0., 0., 0., 0., 1., 0.],
        [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
        [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

print(b)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
        [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
        [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
        [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
        [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
        [0., 0., 1., 1., 1., 1., 1., 1., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

I’m currently using the for loop to do this, are there any ways to do this without using the for loop? Thanks.

  • Only using Pytorch
  • The areas I want to fill always have contiguous boundaries.
  • for loop below loops through row first then column 
b = torch.zeros(a.shape)
for i in range(a.shape[0]):
    occupied = torch.where(a[i] == 1)[0]
    if len(occupied) > 0:
        # indexes in occupied are in ascending order so value of any index
        # between the first and last is filled with 1
        for j in range(occupied[0], occupied[-1] + 1):
            b[i][j] = 1
2

Hi Qimin!

You may use cumsum() to mark those entries that appear to the right
of a 1.0 and similarly use flip() with cumsum() to mark those entries
that appear to the left of a 1.0. You may then use arithmetical (or logical)
operations to combine those pieces together to get your final result:

>>> import torch
>>> print (torch.__version__)
1.13.0
>>>
>>> a = torch.tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
...                   [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
...                   [0., 1., 0., 1., 1., 1., 1., 1., 0., 0.],
...                   [0., 1., 0., 0., 1., 0., 0., 0., 1., 0.],
...                   [0., 1., 0., 1., 0., 0., 0., 0., 1., 0.],
...                   [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
...                   [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
...                   [0., 0., 1., 0., 0., 0., 0., 1., 0., 0.],
...                   [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
...                   [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>>
>>> c = torch.tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
...                   [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
...                   [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
...                   [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
...                   [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
...                   [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
...                   [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
...                   [0., 0., 1., 1., 1., 1., 1., 1., 0., 0.],
...                   [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
...                   [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>>
>>> # using for loop
>>> b = torch.zeros(a.shape)
>>> for i in range(a.shape[0]):
...     occupied = torch.where(a[i] == 1)[0]
...     if len(occupied) > 0:
...         # indexes in occupied are in ascending order so value of any index
...         # between the first and last is filled with 1
...         for j in range(occupied[0], occupied[-1] + 1):
...             b[i][j] = 1
...
>>> torch.equal (b, c)
True
>>>
>>> # loop-free version
>>> aleft = a.flip (1).cumsum (1).flip (1)
>>> aright = a.cumsum (1)
>>> d = ((aleft * aright + a) > 0).float()
>>>
>>> torch.equal (d, c)
True

Best.

K. Frank

3

This is smart! Thank you Frank