test_sparse_tensor.py

test_sparse_tensor.py

import torch
import scipy.sparse as sp
def test_sparse_tensor():
    # That is sparse for sure
    data = np.random.binomial(1, 0.01, size=(128, 27000))
    csr_matrix = sp.csr_matrix(data)

    indptr = csr_matrix.indptr
    indices = csr_matrix.indices
    i = torch.LongTensor(np.array([[l, j] for l, i in enumerate(range(len(indptr) - 1))
                                   for j in indices[indptr[i]:indptr[i + 1]]]))
    sparse_X = torch.sparse.FloatTensor(i.t(), torch.FloatTensor(csr_matrix.data),
                                        torch.Size([int(d) for d in csr_matrix.shape]))
    sparse_X = sparse_X.cuda()

    tic = time.time()
    for i in range(1000):
        sparse_X.to_dense()
    toc = time.time()
    print(toc - tic)

    for i in range(1000):
        torch.FloatTensor(csr_matrix.todense()).cuda()
    tac = time.time()
    print(tac - toc)