柱/行切割火炬稀疏张量

我有一个pytorch稀疏张量,我需要切片行/列使用此切片[idx][:,idx],其中idx是索引列表,使用提到的切片在普通浮点张量上产生我想要的结果.是否可以在稀疏张量上应用相同的切片？这里的例子:

#constructing sparse matrix
i = np.array([[0,1,2,2],[0,1,2,1]])
v = np.ones(4)
i = torch.from_numpy(i.astype("int64"))
v = torch.from_numpy(v.astype("float32"))
test1 = torch.sparse.FloatTensor(i, v)

#constructing float tensor
test2 = np.array([[1,0,0],[0,1,0],[0,1,1]])
test2 = autograd.Variable(torch.cuda.FloatTensor(test2), requires_grad=False)

#slicing
idx = [1,2]
print(test2[idx][:,idx])

输出:

Variable containing:
 1  0
 1  1
[torch.cuda.FloatTensor of size 2x2 (GPU 0)]

我持有250.000 x 250.000邻接矩阵,我需要使用随机idx 对n行和n列进行切片,只需对n随机idx进行采样.由于数据集太大,转换为更方便的数据类型是不现实的.

我可以在test1上实现相同的切片结果吗？它甚至可能吗？如果没有,是否有任何解决方法？

现在我正在使用以下"hack"解决方案来运行我的模型:

idx = sorted(random.sample(range(0, np.shape(test1)[0]), 9000))
test1 = test1AsCsr[idx][:,idx].todense().astype("int32")
test1 = autograd.Variable(torch.cuda.FloatTensor(test1), requires_grad=False)

test1AsCsr是我的test1转换为numpy CSR矩阵的地方.这个解决方案有效,但速度非常慢,并且使我的GPU利用率非常低,因为它需要不断地从CPU内存中读/写.

编辑:结果很好,非稀疏张量