Do a tutorial - sparse COO tensor

Ulf Hamster 4 min.
python pytorch tutorial sparse tensor coo

boilerplate

%%capture 
!pip install torch==1.1.0
# load packages
import torch

# check version
print(torch.__version__)

# set GPU if available
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
1.1.0
cuda:0

Sparse Tensor

pytorch supports sparse tensors in the Coordinate List (COO) format. Each element of a tensor consist of its indicies and the value.

For example, a 2D matrix might have an element $x_{3, 4} = 5.6$. The COO format would store this element as (row=3, col=4, value=5.6) or (3, 4, 5.6).

When implementing COO, indicies and values are usually treated seperatly. First, indicies are usually unsigned integer and values float. Second, it's easier to replace all values without looping over the indicies.

After tensor construction the coordinate list should be ordered by row indices, column indices, and so forth (for multi-dimensional tensors). In pytorch this is done with the .coalesce method.

indices = torch.stack(
    [torch.Tensor([3, 2, 1]), 
     torch.Tensor([4, 2, 0])]).type(torch.uint8)

print(indices)
tensor([[3, 2, 1],
        [4, 2, 0]], dtype=torch.uint8)
weights = torch.Tensor([91, 92, 93])

print(weights)
tensor([91., 92., 93.])
n_dim = 5

matrix = torch.sparse_coo_tensor(
    indices=indices,
    values=weights,
    size=[n_dim, n_dim],
    dtype=torch.uint8,
    device=device
).coalesce()  # "coalesce" orders indicies for faster access

print(matrix)
print("")
print(matrix.to_dense())
tensor(indices=tensor([[1, 2, 3],
                       [0, 2, 4]]),
       values=tensor([93, 92, 91]),
       device='cuda:0', size=(5, 5), nnz=3, dtype=torch.uint8,
       layout=torch.sparse_coo)

tensor([[ 0,  0,  0,  0,  0],
        [93,  0,  0,  0,  0],
        [ 0,  0, 92,  0,  0],
        [ 0,  0,  0,  0, 91],
        [ 0,  0,  0,  0,  0]], device='cuda:0', dtype=torch.uint8)

Example - Sparse Quadratic Matrix

Generate Random Indicies

n_dim = 10
n_cycles = 2
n_weights = n_dim * n_cycles
# reproducability
torch.manual_seed(23)

# initialize lists
indrow = []
indcol = []

# loop n_cycles times 
for _ in range(n_cycles):
    # generate random indicies
    tmp = torch.randperm(n=n_dim)
    # store index
    indrow.append(tmp)
    indcol.append(tmp.roll(1))  # shift it

# stack everything
indices = torch.stack([torch.cat(indrow), torch.cat(indcol)])
print(indices)

# remove double entries, and order COO lists (coalesce)
indices = torch.unique(indices, sorted=True, dim=1)
print(indices)
tensor([[1, 9, 2, 5, 6, 7, 0, 3, 8, 4, 7, 1, 4, 9, 0, 8, 3, 6, 5, 2],
        [4, 1, 9, 2, 5, 6, 7, 0, 3, 8, 2, 7, 1, 4, 9, 0, 8, 3, 6, 5]])
tensor([[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9],
        [7, 9, 4, 7, 5, 9, 0, 8, 1, 8, 2, 6, 3, 5, 2, 6, 0, 3, 1, 4]])

Generate Weights

The elements of our sparse quadratic matrix might be subject to optimization. Thus, we implement them as nn.Parameter.

#weights = torch.nn.Parameter(torch.Tensor(range(n_weights)).to(device))
weights = torch.Tensor(range(n_weights)).to(device)
print(weights)
tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
        14., 15., 16., 17., 18., 19.], device='cuda:0')

Convert to Sparse Matrix

W = torch.sparse_coo_tensor(
    indices=indices, values=weights, 
    size=[10, 10], 
    dtype=torch.float32,
    device=device
).coalesce()

print(W)
print("")
print(W.to_dense())
tensor(indices=tensor([[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9,
                        9],
                       [7, 9, 4, 7, 5, 9, 0, 8, 1, 8, 2, 6, 3, 5, 2, 6, 0, 3, 1,
                        4]]),
       values=tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.,
                      11., 12., 13., 14., 15., 16., 17., 18., 19.]),
       device='cuda:0', size=(10, 10), nnz=20, layout=torch.sparse_coo)

tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.],
        [ 0.,  0.,  0.,  0.,  2.,  0.,  0.,  3.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  4.,  0.,  0.,  0.,  5.],
        [ 6.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  7.,  0.],
        [ 0.,  8.,  0.,  0.,  0.,  0.,  0.,  0.,  9.,  0.],
        [ 0.,  0., 10.,  0.,  0.,  0., 11.,  0.,  0.,  0.],
        [ 0.,  0.,  0., 12.,  0., 13.,  0.,  0.,  0.,  0.],
        [ 0.,  0., 14.,  0.,  0.,  0., 15.,  0.,  0.,  0.],
        [16.,  0.,  0., 17.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0., 18.,  0.,  0., 19.,  0.,  0.,  0.,  0.,  0.]], device='cuda:0')

Apply some (forward) operations

s = torch.ones((n_dim, 1), device=device)
print(s, "\n")

y = torch.matmul(W, s)
print(y, "\n")
tensor([[1.],
        [1.],
        [1.],
        [1.],
        [1.],
        [1.],
        [1.],
        [1.],
        [1.],
        [1.]], device='cuda:0') 

tensor([[ 1.],
        [ 5.],
        [ 9.],
        [13.],
        [17.],
        [21.],
        [25.],
        [29.],
        [33.],
        [37.]], device='cuda:0') 
# NOT SUPPORTED
#print("Let's autograd")
#y.backward()

#print("Check gradients")
#print(weights.grad)

Links