pytorch sparse tensor 在gpu运行时的问题

pytorch 中的稀疏tensor目前只支持coo形式,类似于scipy中coo格式.

eg.

>>> indice = torch.LongTensor([[0, 1, 1],
                          [2, 0, 2]])
>>> value = torch.FloatTensor([3, 4, 5])
>>> torch.sparse.FloatTensor(indice, value, torch.Size([2,3])).to_dense()
 0  0  3
 4  0  5
[torch.FloatTensor of size 2x3]

具体用法可查看文档 https://pytorch.org/docs/stable/sparse.html?highlight=sparse%20tensor

 

当用sparse tensor的代码在gpu上运行时,有个bug,提示indice 和value所在设备应该一致.

RuntimeError: device of indices (0) must match device of values (1)

查看发现它们都是在指定cuda(1)上.

后来在https://github.com/pytorch/pytorch/issues/28500找到解决办法,

torch.sparse_coo_tensor(indices, values.cuda(0)).to(device)

此时代码可以在gpu上跑通了,同时占用2块卡.