我得到了一个带有整数的 pytorch 2-D 张量，并且 2 个整数总是出现在张量的每一行中。
我想创建一个二进制掩码，在这 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].
是否有任何有效且快速的方法来计算此掩码而无需迭代？
【问题讨论】：
标签： python pytorch numpy-ndarray tensor binary-matrix
可能有点混乱，但它无需迭代即可工作。在下文中，我假设一个示例张量 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] 的列索引，以及@对于m 的每一行，987654329@ 类似地包含找到vals[1] 的列索引。
【讨论】：
vals[0] 或 vals[1] 值之一在一行中发生不止一次，这将失败。在编写您提到的行时，我的假设是每个值每行仅出现一次，因此由于 v.nonzero()[:,0] is the row indexes, I could discard this column as I would already know that v.nonzero()[0,:]` 和 v.nonzero()[1:] 对应于 @ 的行 0 987654337@等等。
nonzero()，使用一个返回非零元素的第一次和最后一次出现的函数。
完全基于之前的解决方案，这里是修改后的解决方案：
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)
【讨论】：