使用特定值创建pytorch张量二进制掩码

Question

我得到了一个pytorch 2-D张量，其中有整数，以及总是出现在张量每行中的2个整数。我想创建一个二进制掩码，在这两个整数的两个外观之间包含1，否则为0。例如，如果整数为4和2，并且1-D数组为[1,1,9,4,6,5,1,2,9,9,11,4,3,6,5,2,3,4]，则返回的掩码为：[0,0,0,1,1,1,1,1,0,0,0,1,1,1,1,1,0,0,0].有没有什么高效快捷的方法可以在不迭代的情况下计算这个掩码？

Answer 1

也许有点混乱，但它可以在没有迭代的情况下工作。在下面的例子中，我假设一个张量m，我将其应用于该解，用它来解释比使用通用符号更容易。

import torch
vals=[2,8]#let's assume those are the constant values that appear in each row
#target tensor
m=torch.tensor([[1., 2., 7., 8., 5.],
[4., 7., 2., 1., 8.]])
#let's find the indexes of those values
k=m==vals[0]
p=m==vals[1]
v=(k.int()+p.int()).bool()
nz_indexes=v.nonzero()[:,1].reshape(m.shape[0],2)
#let's create a tiling of the indexes
q=torch.arange(m.shape[1])
q=q.repeat(m.shape[0],1)
#you only need two masks, no matter the size of m. see explanation below
msk_0=(nz_indexes[:,0].repeat(m.shape[1],1).transpose(0,1))<=q
msk_1=(nz_indexes[:,1].repeat(m.shape[1],1).transpose(0,1))>=q
final_mask=msk_0.int() * msk_1.int()
print(final_mask)

我们得到

tensor([[0, 1, 1, 1, 0],
[0, 0, 1, 1, 1]], dtype=torch.int32)

关于两个掩码mask_0和mask_1(如果不清楚它们是什么(，请注意，对于m的每一行，nz_indexes[:,0]包含找到vals[0]的列索引，并且对于m的每一行都，nz_indexes[:,1]类似地包含找到vals[1]的列索引。

Answer 2

完全基于以前的解决方案，下面是经过修订的解决方案：

import torch
vals=[2,8]#let's assume those are the constant values that appear in each row
#target tensor
m=torch.tensor([[1., 2., 7., 8., 5., 2., 6., 5., 8., 4.],
[4., 7., 2., 1., 8., 2., 6., 5., 6., 8.]])
#let's find the indexes of those values
k=m==vals[0]
p=m==vals[1]
v=(k.int()+p.int()).bool()
nz_indexes=v.nonzero()[:,1].reshape(m.shape[0],4)
#let's create a tiling of the indexes
q=torch.arange(m.shape[1])
q=q.repeat(m.shape[0],1)
#you only need two masks, no matter the size of m. see explanation below
msk_0=(nz_indexes[:,0].repeat(m.shape[1],1).transpose(0,1))<=q
msk_1=(nz_indexes[:,1].repeat(m.shape[1],1).transpose(0,1))>=q
msk_2=(nz_indexes[:,2].repeat(m.shape[1],1).transpose(0,1))<=q
msk_3=(nz_indexes[:,3].repeat(m.shape[1],1).transpose(0,1))>=q
final_mask=msk_0.int() * msk_1.int() + msk_2.int() * msk_3.int()
print(final_mask)

我们终于得到

tensor([[0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
[0, 0, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int32)