Pytorch 选择性元件的最大操作，而不是所有元件

我正在给手电筒编码。在torch推断代码之间，为了我自己的兴趣，我添加了一些外围代码。这段代码工作正常，但速度太慢。原因可能是迭代。所以，我需要并行和快速的方法来实现这一点

在张量、Numpy或python数组中这样做是可以的

我创建了一个名为

selective_max

的函数来查找数组中的最大值。但问题是，我不希望在整个数组中有一个最大值，而是在由

mask

array指定的特定候选数组中。让我来说明这个函数的要点（下面是代码本身）

输入

x[batch\u size，dim，num\u points，k]

：x是原始输入，但它通过

x.permute（0,2,1,3）

变成[batch\u size，num\u points，dim，k]

batch_size

是深度学习社会中一个众所周知的定义。在每个小批量中，都有许多点。单个点由

dim

长度特征表示。对于每个特征元素，都有

k

潜在候选，这是

max

函数稍后的目标

mask[批次大小，数量点，k]

：此数组类似于

x

，没有

dim

。其元素为

0

或

1

。所以，我使用它作为掩码信号，就像只对

1

掩码值执行max操作一样

请参阅下面的代码和解释。我使用

3

进行迭代。假设我们针对特定批次和特定点。对于特定批次和特定点，

x

具有[dim，k]数组。掩码具有[k]数组，该数组由

0

或

1

组成。因此，我从[k]数组中提取非零索引，并使用它来提取

x

dim by dim（'for k in range（dim）'）中的特定元素

玩具示例 假设我们在第二个迭代中。因此，我们现在有

[dim，k]

用于

x

和

[k]

用于

掩码。对于这个玩具示例，
i
假定
k=3
和
dim=4
<代码>x=[[3,2,1]，[5,6,4]，[9,8,7]，[12,11,10]

，

k=[0,1,1]

。因此，输出将是

[2,6,8,11]

，而不是

[3,6,9,12]

前科 我尝试

{mask.repeat（0,0,1,0）*（element-wise mul）x}

并执行

max

操作。但是，“0”可能是最大值，因为x在所有数组中可能都有负值。因此，这将导致错误的操作

def selective_max2(x, mask): # x : [batch_size , dim, num_points, k] , mask : [batch_size, num_points, k]
batch_size = x.size(0)
dim = x.size(1)
num_points = x.size(2)
k = x.size(3)
device = torch.device('cuda')

x = x.permute(0,2,1,3) # : [batch, num_points, dim, k]
#print('permuted x dimension : ',x.size())

x = x.detach().cpu().numpy()
mask = mask.cpu().numpy()
output = np.zeros((batch_size,num_points,dim))

for i in range(batch_size):
 for j in range(num_points):
  query=np.nonzero(mask[i][j]) # among mask entries, we get the index of nonzero values.
  for k in range(dim): # for different k values, we get the max value.
   # query is index of nonzero values. so, using query, we can get the values that we want.
   output[i][j][k] = np.max(x[i][j][k][query])

output = torch.from_numpy(output).float().to(device=device)
output = output.permute(0,2,1).contiguous()
return output

免责声明：我已经按照您的玩具示例（但是保留了通用性）编写了以下解决方案

第一件事是（将两者视为）：

然后，选择

1

存在于

k_展开的
中的元素，结果张量为
x
行数（写为
x.shape[0]
），而
k
中的
1
数（或掩码）为列数。到目前为止，我们已经选择了要查询最大元素的范围。然后，使用查找行维度（如中所示）上的最大值

基准

def find_max_elements_inside_tensor_range(arr, mask, return_indices=False):
    mask_expanded = mask.expand_as(arr)
    values, indices = x[k_expanded==1].view(x.shape[0], (k == 1).sum(0)).max(1)
    return (values, indices) if return_indices else values

刚刚添加了第三个参数，以防您想要获得数字索引

%timeit find_max_elements_inside_tensor_range(x, k)
38.4 µs ± 534 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

注：上述解决方案也适用于各种形状的张量和遮罩

def find_max_elements_inside_tensor_range(arr, mask, return_indices=False):
    mask_expanded = mask.expand_as(arr)
    values, indices = x[k_expanded==1].view(x.shape[0], (k == 1).sum(0)).max(1)
    return (values, indices) if return_indices else values

%timeit find_max_elements_inside_tensor_range(x, k)
38.4 µs ± 534 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)