bert-masking.py

bert-masking.py

# Run once for each size.

In [1]: import torch

In [2]: a = torch.arange(100).view(10, 10)

In [3]: a
Out[3]:
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],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
        [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
        [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
        [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
        [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])

In [4]: batch_index = torch.LongTensor([0, 3, 4])

In [5]: batch_index = torch.LongTensor([0, 0, 3])

In [6]: a[batch_index]
Out[6]:
tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]])

In [7]: pos = torch.LongTensor([1, 4, 0])

In [8]: size = torch.LongTensor([3, 2, 4])

In [9]: a[batch_index]
Out[9]:
tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]])

In [10]: a[batch_index, pos]
Out[10]: tensor([ 1,  4, 30])

In [11]: a[batch_index, pos:pos+size]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-a83bf9a7324d> in <module>()
----> 1 a[batch_index, pos:pos+size]

TypeError: only integer tensors of a single element can be converted to an index

In [12]: mask = torch.LongTensor(*a[batch_index].shape).fill_(0)

In [13]: mask
Out[13]:
tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

In [14]: for i in range(3):
    ...:     mask[pos[i]:pos[i]+size[i]] = 1
    ...:

In [15]: mask
Out[15]:
tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

In [16]: mask = torch.LongTensor(*a[batch_index].shape).fill_(0)

In [17]: for i in range(3):
    ...:     mask[i, pos[i]:pos[i]+size[i]] = 1
    ...:
    ...:

In [18]: mask
Out[18]:
tensor([[0, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0, 0]])

In [19]: a[batch_index]
Out[19]:
tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]])

In [20]: a[batch_index][mask == 1]
Out[20]: tensor([ 1,  2,  3,  4,  5, 30, 31, 32, 33])

In [21]: size = torch.LongTensor([3, 3, 3])

In [22]: mask = torch.LongTensor(*a[batch_index].shape).fill_(0)

In [23]: for i in range(3):
    ...:     mask[i, pos[i]:pos[i]+size[i]] = 1
    ...:
    ...:

In [24]: mask
Out[24]:
tensor([[0, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0]])

In [25]: a[batch_index][mask == 1].view(3, -1)
Out[25]:
tensor([[ 1,  2,  3],
        [ 4,  5,  6],
        [30, 31, 32]])