Linear Algebra in PyTorch
Contents
Linear Algebra in PyTorch¶
PyTorch is a popular package for developing models for deep learning. In this section, we’ll look at its linear algebra capabilities.
Even if you are not doing deep learning, you can use PyTorch for linear algebra. One of the nice things about PyTorch is that it makes it easy to take advantage of GPU hardware, which is very efficient at certain operations.
You can find PyTorch tutorials on the official website as well as documentation. Note these tutorials are focused on deep learning. This section will focus on the linear algebra capabilities.
When you install PyTorch, use the pytorch
channel in conda.
PyTorch Tensors¶
PyTorch Tensor
s are just multi-dimensional arrays. You can go back and forth between these and numpy ndarray
.
To put a tensor on gpu, use cuda()
. Note that you must have an NVIDIA GPU in your computer to be able to do this successfully.
To move a tensor back to CPU, you can use device='cpu'
Linear Algebra Functions¶
PyTorch provides access to a variety of BLAS and LAPACK-type routines - see documentation here. These do not follow the BLAS/LAPACK naming conventions
torch.addmv
is roughly equivalent to axpy
, and performs
Let’s look at the timing difference between CPU and GPU
m = 1000
n = 1000
Anp = np.random.randn(m,n)
xnp = np.random.randn(n)
ynp = np.random.randn(m)
print("numpy")
%time z = ynp + Anp @ xnp
for device in ('cpu', 'cuda'):
print(f"\ndevice = {device}")
A = torch.Tensor(Anp).to(device=device)
x = torch.Tensor(xnp).to(device=device)
y = torch.Tensor(ynp).to(device=device)
z = torch.addmv(y, A, x)
%time z = torch.addmv(y, A, x)
torch.mv
performs matrix-vector products
torch.mm performs matrix-matrix multiplications
m = 1000
n = 1000
Anp = np.random.randn(m,n)
Bnp = np.random.randn(n, n)
print("numpy")
%time C = Anp @ Bnp
for device in ['cpu', 'cuda']:
print(f"\ndevice = {device}")
A = torch.Tensor(Anp).to(device=device)
B = torch.Tensor(Bnp).to(device=device)
C = torch.mm(A, B) # run once to warm up
%time C = torch.mm(A, B)
Batch operations¶
Where PyTorch (and GPUs in general) really shine are in batch operations. We get extra efficiency if we do a bunch of multiplications with matrices of the same size.
For matrix-matrix multiplcation, the function is torch.bmm
Because tensors are row-major, we want the batch index to be the first index. In the below code, the batch multiplication is equivalent to
n = 512 # matrix size
k = 32 # batch size
Anp = np.random.randn(k, n, n)
Bnp = np.random.randn(k, n, n)
# see numpy matmul documentation for how this performs batch multiplication
print("numpy")
%time C = np.matmul(Anp, Bnp)
for device in ['cpu', 'cuda']:
print(f"\ndevice = {device}")
A = torch.randn(k, n, n).to(device=device)
B = torch.randn(k, n, n).to(device=device)
C = torch.bmm(A, B) # run once to warm up
%time C = torch.bmm(A, B)
Sparse Linear Algebra¶
PyTorch also supports sparse tensors in torch.sparse
. Tensors are stored in COOrdinate format.
indices are stored in a 2 x nnz
tensor of Long
(a datatype that stores integers). Values are stored as floats.
Exercise¶
Write a function that returns a sparse identity matrix of size n
in PyTorch.