向量化索引识别以加快 GPU 处理速度

Question

我的输入是一个值列表，data_idx。在示例中，值的范围为 [0, 5]。

data_idx = [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]

我想要的输出 filled_matrix 是一个形状为 max(value) 的张量 len(data_idx) ，其中张量的每一行 r 包含所有索引，其中 data_idx == r 和 -1 表示该行的其余部分如果匹配索引的数量少于len(data_idx)

例如，在第一行中，r=0、data_idx==0 在索引 [3, 11, 14] 处。完整的输出如下所示：

filled_matrix = tensor([[ 3, 11, 14, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1], 
    [ 5, 10, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 0,  9, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 8, 12, 13, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 4,  6, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 1,  2,  7, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1]],
   dtype=torch.int8)

我有可以实现我的目标的循环代码。

import torch

max_idx = 6
data_idx = torch.tensor([2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]).cuda()

max_number_data_idx = data_idx.shape[0]

filled_matrix = torch.zeros([max_idx, max_number_data_idx], dtype=torch.int8, device='cuda')
filled_matrix.fill_(-1)
for i in range(max_idx):
   same_idx = (data_idx == i).nonzero().flatten()
   filled_matrix[i][:same_idx.shape[0]] = same_idx

现在，我想加快这段代码的速度。具体来说，我希望它在 GPU 上更快。在实际场景中，输入data_idx 可以是包含数百万个值的列表。在这种情况下，例如 1 M 的不同值，GPU 将被调用 1 M 的时间，这使得它非常慢。我的代码是顺序的，GPU 讨厌顺序代码。

是否有一个函数可以更有效地产生相同的结果？或者一种向量化这个 for 循环的方法？

Answer 1

免责声明：我还没有分析过这段代码，看看它在 GPU 上是否真的更快。

一种矢量化解决方案是使用张量视图来广播比较。张量视图不使用额外的内存。更多详情请看documentation

首先，创建一个矩阵，其中包含您要比较的每一行的值。在这种情况下，它只是行索引。

comparison = torch.tensor(range(max_idx))

现在我们将使用expand 和unsqueeze 来创建与filled_matrix 形状相同的data_idx 和comparison 的视图。

comparison_view = comparison.unsqueeze(1).expand(max_idx, max_number_data_idx)
print(comparison_view)

# Each row is the index you want to compare to
# tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
    [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4],
    [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5]])

data_idx_view = data_idx.expand(max_idx, max_number_data_idx)
print(data_idx_view)

# Each row is a copy of data_idx
# tensor([[2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]])

我们可以比较它们的相等性并使用nonzero 来查找索引

mask = comparison_view == data_idx_view
mask_indices = mask.nonzero()

print(mask_indices)
# tensor([[ 0,  3],
    [ 0, 11],
    [ 0, 14],
    [ 1,  5],
    [ 1, 10],
    [ 2,  0],
    [ 2,  9],
    [ 3,  8],
    [ 3, 12],
    [ 3, 13],
    [ 4,  4],
    [ 4,  6],
    [ 5,  1],
    [ 5,  2],
    [ 5,  7]])

现在，您只需将这些结果处理成您想要的输出格式。

filled_matrix = torch.zeros([max_idx, max_number_data_idx], dtype=torch.int8)
filled_matrix.fill_(-1)
col_indices = [0, 1, 2, 0, 1, 0, 1, 0, 1, 2, 0, 1, 0, 1, 2]
filled_matrix[mask_indices[:, 0], col_indices] = mask_indices[:, 1].type(torch.int8)

我考虑了几种生成col_indices 列表的选项，但如果没有for 循环，我什么都想不出来。

col_indices = torch.zeros(mask_indices.shape[0])
for i in range(1, mask_indices.shape[0]):
    if mask_indices[i,0] == mask_indices[i-1,0]:
        col_indices[i] = col_indices[i-1]+1

您需要进行一些分析以查看哪些代码实际上更快。