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

Question

我得到了一个带有整数的 pytorch 2-D 张量，以及始终出现在张量的每一行中的 2 个整数。我想创建一个二进制掩码，在这两个整数的两次出现之间包含 1，否则为 0。例如，如果整数是 4 和 2，并且一维数组是[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]. Is there any无需迭代即可计算此掩码的有效且快速的方法？

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如果不清楚它们是什么，请注意包含，对于 的每一nz_indexes[:,0]行，找到m的列索引vals[0]，并且nz_indexes[:,1]类似地包含，对于 的每一行m，找到的列索引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)