Do a tutorial - Random Sparse tensor

Ulf Hamster 7 min.
python pytorch tutorial sparse tensor coo random randn_coo_matrix apply_to_nonzero

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

Generate Random Sparse COO Matrix

def randn_coo_matrix(n_rows, n_cols, n_elem, 
                     dtype=torch.float32, 
                     device=torch.device("cpu"),
                     random_state=None):
    # set seed
    if random_state:
        torch.random.manual_seed(random_state)
    
    # generate random indicies
    indices = torch.stack([
        torch.randint(size=(n_elem,), high=n_rows),
        torch.randint(size=(n_elem,), high=n_cols)
    ]).to(device)
    
    # coalesce
    indices = torch.unique(indices, sorted=True, dim=1)
    n_elem2 = indices.size()[1]
    
    # generate normal distributed values
    values = torch.randn((n_elem2,), device=device)
    
    # create COO matrix
    return torch.sparse_coo_tensor(
        indices=indices, 
        values=values, 
        size=[n_rows, n_cols], 
        dtype=dtype,
        device=device
    ).coalesce()

X = randn_coo_matrix(5000, 10000, 30000, device=device, random_state=42)
X
tensor(indices=tensor([[   0,    0,    0,  ..., 4999, 4999, 4999],
                       [ 171, 1025, 2556,  ..., 6345, 6979, 7666]]),
       values=tensor([ 0.6226,  0.8774,  1.1547,  ...,  1.3963, -0.6628,
                       0.6630]),
       device='cuda:0', size=(5000, 10000), nnz=29991, layout=torch.sparse_coo)

Basic Methods

Tensor Properties

torch.manual_seed(42)
A = randn_coo_matrix(3, 4, 7, device=device)
print(A.to_dense())
tensor([[ 0.6226,  0.0000,  0.8774,  0.0000],
        [ 0.0000,  1.1547,  1.3588,  0.0000],
        [ 0.0000,  0.0000, -0.8027,  0.0000]], device='cuda:0')

matrix dimensions

A.size()
torch.Size([3, 4])
n_rows, n_cols = A.size()
n_rows, n_cols
(3, 4)

number of non-zero elements

A._nnz()
5

indices of non-zero-elements (COO)

A.indices()
tensor([[0, 0, 1, 1, 2],
        [0, 2, 1, 2, 2]], device='cuda:0')

corresponding element values (COO)

A.values()
tensor([ 0.6226,  0.8774,  1.1547,  1.3588, -0.8027], device='cuda:0')

coalesce status (COO)

A.is_coalesced()
True
A = A.coalesce()

data type

A.dtype
torch.float32

CPU vs GPU

A.device
device(type='cuda', index=0)
A.get_device()
0

Reset Sparse Matrix to zeros

torch.manual_seed(42)
A = randn_coo_matrix(3, 4, 7, device=device)
print(A)
tensor(indices=tensor([[0, 0, 1, 1, 2],
                       [0, 2, 1, 2, 2]]),
       values=tensor([ 0.6226,  0.8774,  1.1547,  1.3588, -0.8027]),
       device='cuda:0', size=(3, 4), nnz=5, layout=torch.sparse_coo)
A.zero_()
print(A)
tensor(indices=tensor([], size=(2, 0)),
       values=tensor([], size=(0,)),
       device='cuda:0', size=(3, 4), nnz=0, layout=torch.sparse_coo)

Different Variables, Same Data

First, create a random sparse tensor A

torch.manual_seed(42)
A = randn_coo_matrix(2, 3, 4, device=device)
print(A.to_dense())
tensor([[0.6226, 0.8774, 1.1547],
        [0.0000, 0.0000, 1.3588]], device='cuda:0')

Assign A to Z. Manipulate an element in Z.

Z = A
Z.values()[0] = 999

The change will be in Z as well as a A

print(Z.to_dense())
print(A.to_dense())
tensor([[9.9900e+02, 8.7742e-01, 1.1547e+00],
        [0.0000e+00, 0.0000e+00, 1.3588e+00]], device='cuda:0')
tensor([[9.9900e+02, 8.7742e-01, 1.1547e+00],
        [0.0000e+00, 0.0000e+00, 1.3588e+00]], device='cuda:0')

Deep Copy – clone

Again, generate a random sparse tensor A

torch.manual_seed(42)
A = randn_coo_matrix(2, 3, 4, device=device)
print(A.to_dense())
tensor([[0.6226, 0.8774, 1.1547],
        [0.0000, 0.0000, 1.3588]], device='cuda:0')

Now call the .clone() method to copy all the data from A to Z. And manipulate an element Z again.

Z = A.clone()
Z.values()[0] = 999

This time A is not affected because Z and A point to different data in the memory.

print(Z.to_dense())
print(A.to_dense())
tensor([[9.9900e+02, 8.7742e-01, 1.1547e+00],
        [0.0000e+00, 0.0000e+00, 1.3588e+00]], device='cuda:0')
tensor([[0.6226, 0.8774, 1.1547],
        [0.0000, 0.0000, 1.3588]], device='cuda:0')

Slice – narrow_copy

The method .narrow_copy(..) will copy a range of rows (dim=0) or range of columns (dim=1) into a new sparse tensor.

torch.manual_seed(42)
A = randn_coo_matrix(5, 7, 30, device=device)
print(A.to_dense())
tensor([[ 0.6226,  0.0000,  0.8774,  0.0000,  1.1547,  1.3588, -0.8027],
        [ 0.7622,  0.0000, -2.0681,  0.0000,  0.0000,  1.6839, -0.8694],
        [ 0.0000,  0.0000, -0.1916, -0.4707,  0.8071, -0.1908,  0.0000],
        [-1.4708,  0.0930,  0.0000,  0.0000, -0.7613, -0.1699,  0.0000],
        [ 0.6272, -0.7603,  0.0000,  0.9738,  0.0000, -0.4690,  0.0000]],
       device='cuda:0')

Copy the the second to third row of A into Z

Z = A.narrow_copy(dim=0, start=1, length=2)
print(Z.to_dense())
tensor([[ 0.7622,  0.0000, -2.0681,  0.0000,  0.0000,  1.6839, -0.8694],
        [ 0.0000,  0.0000, -0.1916, -0.4707,  0.8071, -0.1908,  0.0000]],
       device='cuda:0')

Copy the third to fifth column from A to Z

Z = A.narrow_copy(dim=1, start=2, length=3)
print(Z.to_dense())
tensor([[ 0.8774,  0.0000,  1.1547],
        [-2.0681,  0.0000,  0.0000],
        [-0.1916, -0.4707,  0.8071],
        [ 0.0000,  0.0000, -0.7613],
        [ 0.0000,  0.9738,  0.0000]], device='cuda:0')

Combine both slices

Z = A.narrow_copy(dim=0, start=1, length=2).narrow_copy(dim=1, start=2, length=3)
print(Z.to_dense())
tensor([[-2.0681,  0.0000,  0.0000],
        [-0.1916, -0.4707,  0.8071]], device='cuda:0')

Math methods

torch.manual_seed(42)
A = randn_coo_matrix(3, 5, 7, device=device)
B = randn_coo_matrix(3, 5, 7, device=device)
print(A.to_dense())
print(B.to_dense())
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.6226],
        [0.0000, 0.0000, 0.0000, 0.8774, 1.1547],
        [1.3588, 0.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')
tensor([[ 0.0000,  1.4505,  0.0000,  0.4870,  0.0000],
        [ 0.0000,  1.0220,  0.0000,  1.2476, -0.8966],
        [ 0.0000, -0.0573,  0.0000,  0.0000,  0.0000]], device='cuda:0')

Addition

Z = A.add(B)
print(Z.to_dense())
tensor([[ 0.0000,  1.4505,  0.0000,  0.4870,  0.6226],
        [ 0.0000,  1.0220,  0.0000,  2.1251,  0.2581],
        [ 1.3588, -0.0573,  0.0000,  0.0000,  0.0000]], device='cuda:0')
Z = A + B
print(Z.to_dense())
tensor([[ 0.0000,  1.4505,  0.0000,  0.4870,  0.6226],
        [ 0.0000,  1.0220,  0.0000,  2.1251,  0.2581],
        [ 1.3588, -0.0573,  0.0000,  0.0000,  0.0000]], device='cuda:0')

Subtract

Z = A.sub(B)
print(Z.to_dense())
tensor([[ 0.0000, -1.4505,  0.0000, -0.4870,  0.6226],
        [ 0.0000, -1.0220,  0.0000, -0.3702,  2.0513],
        [ 1.3588,  0.0573,  0.0000,  0.0000,  0.0000]], device='cuda:0')
Z = A - B
print(Z.to_dense())
tensor([[ 0.0000, -1.4505,  0.0000, -0.4870,  0.6226],
        [ 0.0000, -1.0220,  0.0000, -0.3702,  2.0513],
        [ 1.3588,  0.0573,  0.0000,  0.0000,  0.0000]], device='cuda:0')

Divide by 1 scalar

Z = A.div(2.34)
print(Z.to_dense())
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.2661],
        [0.0000, 0.0000, 0.0000, 0.3750, 0.4934],
        [0.5807, 0.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')
Z = A / 2.34
print(Z.to_dense())
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.2661],
        [0.0000, 0.0000, 0.0000, 0.3750, 0.4934],
        [0.5807, 0.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')

Multiply elementwise

Z = A.mul(B)
print(Z.to_dense())
tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.0947, -1.0353],
        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]], device='cuda:0')
Z = A * B
print(Z.to_dense())
tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.0947, -1.0353],
        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]], device='cuda:0')

Elementwise multiplication of sparse matrices usually lead to an even sparser matrix due to multiplications with zeros.

print(Z)
tensor(indices=tensor([[1, 1],
                       [3, 4]]),
       values=tensor([ 1.0947, -1.0353]),
       device='cuda:0', size=(3, 5), nnz=2, layout=torch.sparse_coo)

Sparse <-> Dense

torch.manual_seed(42)
A = randn_coo_matrix(3, 2, 4, device=device)  # Sparse Matrix
B = torch.randn(3, 2, device=device)  # Dense Matrix
A
tensor(indices=tensor([[0, 1, 2],
                       [0, 0, 1]]),
       values=tensor([0.6226, 0.8774, 1.1547]),
       device='cuda:0', size=(3, 2), nnz=3, layout=torch.sparse_coo)

Sparse and dense tensors are usually incompatible.

try:
    Z = A + B
except:
    print("doesn't work")
doesn't work

Convert the sparse matrix to dense.

Z = A.to_dense() + B
print(Z)
tensor([[ 2.0731,  0.4870],
        [ 1.8994,  1.2476],
        [-0.8966,  1.0974]], device='cuda:0')

Or the dense matrix to a sparse matrix

Z = A + B.to_sparse()
print(Z.to_dense())
tensor([[ 2.0731,  0.4870],
        [ 1.8994,  1.2476],
        [-0.8966,  1.0974]], device='cuda:0')

Functions

Sum

torch.manual_seed(42)
A = randn_coo_matrix(3, 5, 7, device=device)
print(A.to_dense())
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.6226],
        [0.0000, 0.0000, 0.0000, 0.8774, 1.1547],
        [1.3588, 0.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')

Sum up each column (dim=0)

col_sum = torch.sparse.sum(A, dim=0)
print(col_sum.to_dense())
tensor([1.3588, 0.0000, 0.0000, 0.8774, 1.7773], device='cuda:0')

Sum up each row (dim=1)

row_sum = torch.sparse.sum(A, dim=1)
print(row_sum.to_dense())
tensor([0.6226, 2.0321, 1.3588], device='cuda:0')

Multiply Sparse Matrix with Dense Matrix – mm

$Z = A \cdot B$ with

torch.manual_seed(42)
A = randn_coo_matrix(3, 5, 7, device=device)  # Sparse Matrix
B = torch.randn(5, 2, device=device)  # Dense Matrix
Z = torch.sparse.mm(A, B)
print(Z)  # Dense Matrix
tensor([[-0.5475, -0.0169],
        [-3.8243,  0.6395],
        [ 1.9709,  0.6617]], device='cuda:0')

addmm

$Z = \alpha \cdot (A \cdot B) + \beta \cdot C$ with

torch.manual_seed(42)
A = randn_coo_matrix(3, 5, 7, device=device)  # Sparse Matrix
B = torch.randn(5, 2, device=device)  # Dense Matrix
C = torch.randn(3, 2, device=device)  # Dense Matrix
alp = 1
bet = 1
Z = torch.sparse.addmm(C, A, B, alpha=alp, beta=bet)
print(Z)
tensor([[ 0.0848,  0.5624],
        [-4.5831,  1.7189],
        [ 0.5028,  1.7162]], device='cuda:0')

Transformation of non-zero elements

torch.sparse does not implement simple transformations, e.g. log, exp, abs, etc. There are two ways to implement such transformations

Example – abs

torch.manual_seed(42)
A = randn_coo_matrix(3, 4, 9, device=device)  # Sparse Matrix
print(A.to_dense())
tensor([[ 0.0000,  0.6226,  0.8774,  0.0000],
        [ 1.1547,  0.0000,  1.3588, -0.8027],
        [ 0.0000,  0.0000,  0.7622, -2.0681]], device='cuda:0')
fn = torch.abs
B = torch.sparse_coo_tensor(
    indices=A.indices(), 
    values=fn(A.values()), 
    size=A.size(), 
    dtype=A.dtype,
    device=A.device)
print(B.to_dense())
tensor([[0.0000, 0.6226, 0.8774, 0.0000],
        [1.1547, 0.0000, 1.3588, 0.8027],
        [0.0000, 0.0000, 0.7622, 2.0681]], device='cuda:0')

apply_to_nonzero

def apply_to_nonzero(x, fn=torch.Tensor):
    return torch.sparse_coo_tensor(
        indices=x.indices(), 
        values=fn(x.values()), 
        size=x.size(), 
        dtype=x.dtype,
        device=x.device).coalesce()
Z = apply_to_nonzero(A, fn=torch.exp)
print(Z.to_dense())
tensor([[0.0000, 1.8638, 2.4047, 0.0000],
        [3.1730, 0.0000, 3.8914, 0.4481],
        [0.0000, 0.0000, 2.1429, 0.1264]], device='cuda:0')
Z = apply_to_nonzero(A, fn=torch.tanh)
print(Z.to_dense())
tensor([[ 0.0000,  0.5529,  0.7051,  0.0000],
        [ 0.8193,  0.0000,  0.8761, -0.6655],
        [ 0.0000,  0.0000,  0.6423, -0.9685]], device='cuda:0')
Z = apply_to_nonzero(A, fn=torch.abs)
print(Z.to_dense())
tensor([[0.0000, 0.6226, 0.8774, 0.0000],
        [1.1547, 0.0000, 1.3588, 0.8027],
        [0.0000, 0.0000, 0.7622, 2.0681]], device='cuda:0')
Z = apply_to_nonzero(A, fn=torch.relu)
print(Z)  # 
tensor(indices=tensor([[0, 0, 1, 1, 1, 2, 2],
                       [1, 2, 0, 2, 3, 2, 3]]),
       values=tensor([0.6226, 0.8774, 1.1547, 1.3588, 0.0000, 0.7622, 0.0000]),
       device='cuda:0', size=(3, 4), nnz=7, layout=torch.sparse_coo)
Z = apply_to_nonzero(A, fn=lambda e: e+10)
print(Z.to_dense())
tensor([[ 0.0000, 10.6226, 10.8774,  0.0000],
        [11.1547,  0.0000, 11.3588,  9.1973],
        [ 0.0000,  0.0000, 10.7622,  7.9319]], device='cuda:0')
Z = apply_to_nonzero(A, fn=lambda e: torch.log(torch.abs(e)))
print(Z.to_dense())
tensor([[ 0.0000, -0.4738, -0.1308,  0.0000],
        [ 0.1438,  0.0000,  0.3066, -0.2198],
        [ 0.0000,  0.0000, -0.2716,  0.7266]], device='cuda:0')

Links