柱/行切割火炬稀疏张量

Question

我有一个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内存中读/写.

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

Answer 1

好吧，这个问题已经有几年了，但迟到总比不到好。

\n

这是我用来切片稀疏张量的函数。（辅助函数如下）

\n

def slice_torch_sparse_coo_tensor(t, slices):\n    """\n    params:\n    -------\n    t: tensor to slice\n    slices: slice for each dimension\n\n    returns:\n    --------\n    t[slices[0], slices[1], ..., slices[n]]\n    """\n\n    t = t.coalesce()\n    assert len(args) == len(t.size())\n    for i in range(len(args)):\n        if type(args[i]) is not torch.Tensor:\n            args[i] = torch.tensor(args[i], dtype= torch.long)\n\n    indices = t.indices()\n    values = t.values()\n    for dim, slice in enumerate(args):\n        invert = False\n        if t.size(0) * 0.6 < len(slice):\n            invert = True\n            all_nodes = torch.arange(t.size(0))\n            unique, counts = torch.cat([all_nodes, slice]).unique(return_counts=True)\n            slice = unique[counts==1]\n        if slice.size(0) > 400:\n            mask = ainb_wrapper(indices[dim], slice)\n        else:\n            mask = ainb(indices[dim], slice)\n        if invert:\n            mask = ~mask\n        indices = indices[:, mask]\n        values = values[mask]\n\n        \n    return torch.sparse_coo_tensor(indices, values, t.size()).coalesce()\n

\n

使用情况（在我的机器上花了 2.4 秒）：

\n

indices = torch.randint(low= 0, high= 200000, size= (2, 1000000))\nvalues = torch.rand(size=(1000000,))\nt = torch.sparse_coo_tensor(indices, values, size=(200000, 200000))\nidx = torch.arange(1000)\nslice_coo(t, [idx, idx])\n

\n

出去：

\n

tensor(indices=tensor([[ 13,  62,  66,  78, 134, 226, 233, 266, 299, 344, 349,\n                        349, 369, 396, 421, 531, 614, 619, 658, 687, 769, 792,\n                        810, 840, 926, 979],\n                       [255, 479, 305, 687, 672, 867, 444, 559, 772,  96, 788,\n                        980, 423, 699, 911, 156, 267, 721, 381, 781,  97, 271,\n                        840, 292, 487, 185]]),\n       values=tensor([0.4260, 0.4816, 0.8001, 0.8815, 0.3971, 0.4914, 0.7068,\n                      0.2329, 0.4038, 0.1757, 0.7758, 0.3210, 0.2593, 0.8290,\n                      0.1320, 0.4322, 0.7529, 0.8341, 0.8128, 0.4457, 0.4100,\n                      0.1618, 0.4097, 0.3088, 0.6942, 0.5620]),\n       size=(200000, 200000), nnz=26, layout=torch.sparse_coo)\n

\n

slice_torch_sparse_coo_tensor 的计时：

\n

%timeit slice_torch_sparse_coo_tensor(t, torch.randperm(200000)[:500], torch.arange(200000))\n\noutput:\n    1.08 s \xc2\xb1 447 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n

\n

对于内置的 torch.index_select （在此处实现）：

\n

%timeit t.index_select(0, torch.arange(100))\n\noutput:\n    56.7 s \xc2\xb1 4.87 s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n

\n

这些是我用于此目的的辅助函数，函数“ainb”查找 a 中 b 中的元素。我不久前在互联网上发现了这个功能，但我找不到链接它的帖子。

\n

import torch\ndef ainb(a,b):\n    """gets mask for elements of a in b"""\n\n    size = (b.size(0), a.size(0))\n\n    if size[0] == 0: # Prevents error in torch.Tensor.max(dim=0)\n        return torch.tensor([False]*a.size(0), dtype= torch.bool)\n        \n    a = a.expand((size[0], size[1]))\n    b = b.expand((size[1], size[0])).T\n\n    mask = a.eq(b).max(dim= 0).values\n\n    return mask\n\ndef ainb_wrapper(a, b, splits = .72):\n    inds = int(len(a)**splits)\n\n    tmp = [ainb(a[i*inds:(i+1)*inds], b) for i in list(range(inds))]\n\n    return torch.cat(tmp)\n

\n

由于该函数随元素数量呈二次方缩放，因此我添加了一个包装器，将输入分割成块，然后连接输出。仅使用 CPU 效率更高，但我不确定使用 GPU 时这是否成立，如果有人可以测试它，我将不胜感激:)

\n

这是我第一次发帖，因此也感谢有关帖子质量的反馈。

\n