在传播梯度时构造稀疏张量？

Question

我有与此类似的代码，我想做得更快：

# indices: indices of a 3d tensor
# values associated to the indices


result = torch.zeros((L, N, N))
for idx, (i,j,k) in enumerate(indices):
        mask = torch.zeros_like(result)
        mask[i][j][k] = 1.0
        img = img + mask * values[idx]

现在，即使我选择制作稀疏掩模，我也注意到每次迭代运行速度都变慢。对于传播

img = func(indices, values) 

形式的梯度的函数，是否有一个简单的解决方案

我正在寻找一种利用矢量化或/和稀疏数据结构的解决方案

Answer 1

解决方案就像构造一个新的稀疏 coo 张量一样简单。

value_tensor = torch.sparse_coo_tensor(indices=indices, values=values, size=(L, N, N))