PyTorch 稀疏张量上的快速 GPU 计算

心中带点小风骚 2年前 pytorch 416

原文标题 ：Fast GPU computation on PyTorch sparse tensor

是否可以对 PyTorch MxN 张量的每一行进行操作，但只能在某些索引（例如非零）处进行操作以节省时间？

我对 M 和 N 非常大的情况特别感兴趣，其中每行上只有几个元素不为空。

（玩具示例）从这个大张量：

Large = Tensor([[0, 1, 3, 0, 0, 0],
                [0, 0, 0, 0, 5, 0],
                [1, 0, 0, 5, 0, 1]])

我想使用类似以下较小的“张量”：

irregular_tensor = [ [1, 3],
                     [5],
                     [1, 5, 1]]

并对每一行进行相同的精确计算（例如涉及torch.cumsum和torch.exp）以获得大小为 Mx1 的输出。

有没有办法做到这一点？

原文链接：https://stackoverflow.com//questions/71531822/fast-gpu-computation-on-pytorch-sparse-tensor

我来回复

D.L 评论

这用list comprehension方法减少了张量（list of lists），速度很快。

#  original large tensor
Large_tensor = [ [0, 1, 3, 0, 0, 0],
                 [0, 0, 0, 0, 5, 0],
                 [1, 0, 0, 5, 0, 1]]

# size of tensor
imax = len(Large_tensor)
jmax = len(Large_tensor[0])
print(f'the tensor is of size {imax} x {jmax}')

small_tensor = [[x for x in row if x!=0] for row in Large_tensor]

# result
print('large tensor: ', Large_tensor)
print('small tensor: ', small_tensor)

结果是：

the tensor is of size 3 x 6
large tensor:  [[0, 1, 3, 0, 0, 0], [0, 0, 0, 0, 5, 0], [1, 0, 0, 5, 0, 1]]
small tensor:  [[1, 3], [5], [1, 5, 1]]

另一种方法通过iterating通过large_tensor的分量减少张量以创建具有非零值的small_tensor（根据问题）。

#  original large tensor
Large_tensor = [ [0, 1, 3, 0, 0, 0],
                 [0, 0, 0, 0, 5, 0],
                 [1, 0, 0, 5, 0, 1]]

# example of how to reference tensor
i = 0
j = 2
imax = len(Large_tensor)
jmax = len(Large_tensor[0])
print(Large_tensor[i][j])

# the dimension of the tensor
print(f'the tensor is of size {imax} x {jmax}')

# empty list for the new small tensor
small_tensor = []

# process of reducing
for i in range(imax):
    small_tensor.append([])
    for j in range(jmax):
        if Large_tensor[i][j]!=0:
            small_tensor[i].append(Large_tensor[i][j])
            print(i, j)

# result
print('large tensor: ', Large_tensor)
print('small tensor: ', small_tensor)

输出是：

large tensor:  [[0, 1, 3, 0, 0, 0], [0, 0, 0, 0, 5, 0], [1, 0, 0, 5, 0, 1]]
small tensor:  [[1, 3], [5], [1, 5, 1]]

结论：显示了两种快速方法。可能还有更快的方法使用外部模块，例如numpy和SciPy。

2年前 0条评论