What is Torch Vectorized ?¶

Torch Vectorized offers batched and vectorized operations on volume of 3x3 symmetric matrices with Pytorch and their associated differentiable torch.nn.Module layers. The current Pytorch’s implementation of batch eigen-decomposition is very slow when dealing with huge number of small matrices (e.g. 500k x 3x3). This library offers some basic functions like vSymEig, vExpm and vLogm for fast computation (>250x faster) of huge number of small matrices with Pytorch using an analytical solution.

See Ionescu et al., Matrix backpropagation for deep networks with structured layers, CVPR 2015 for details on the gradients computation

vSymEig¶

A quick closed-form solution for volumetric 3x3 matrices Eigen-Decomposition with Pytorch. Solves Eigen-Decomposition of data with shape Bx9xDxHxW, where B is the batch size, 9 is the flattened 3x3 symmetric matrices, D is the depth, H is the Height, W is the width. The goal is to accelerate the Eigen-Decomposition of multiple (>500k) small matrices (3x3) on GPU with Pytorch using an analytical solution.

vExpm¶

Based on vSymEig, computes the matrix exponential for batch of volumetric 3x3 matrices.

vLogm¶

Based on vSymEig, computes the matrix logarithm for batch of volumetric 3x3 matrices.

Install me¶

pip install torch-vectorized

How to use¶

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vSymEig

# Random batch of volumetric 3x3 symmetric matrices of size 16x9x32x32x32
input = sym(torch.rand(16, 9, 32, 32, 32))

# Output eig_vals with size: 16x3x32x32x32 and eig_vecs with size 16,3,3,32,32,32
eig_vals, eig_vecs = vSymEig(input, eigen_vectors=True)

vlinalg¶

torchvectorized.vlinalg.vSymEig(inputs: torch.Tensor, eigenvectors=False, flatten_output=False, descending_eigenvals=False)¶

Compute the eigendecomposition $\mathbf{M} = \mathbf{U} \mathbf{\Sigma} \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices. eigenvectors (bool) – If `True`, computes the eigenvectors. flatten_output (bool) – If `True` the eigenvalues are returned as: *(BDHW)x3 and the eigenvectors as (BDHW)x3x3* otherwise they are returned with shapes Bx3xDxHxW and Bx3x3xDxHxW respectively. descending_eigenvals (bool) – If `True`, return the eigenvvalues in descending order
Returns:	Return the eigenvalues and the eigenvectors as tensors.
Return type:	tuple[torch.Tensor, None]

Example:

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vSymEig

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))
eig_vals, eig_vecs = vSymEig(inputs, eigenvectors=True)

torchvectorized.vlinalg.vExpm(inputs: torch.Tensor, replace_nans=False)¶

Compute the matrix exponential $\mathbf{M} = \mathbf{U} exp(\mathbf{\Sigma}) \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices. replace_nans (bool) – If `True`, replace nans by 0
Returns:	Return a tensor with shape Bx9xDxHxW where every voxel is the matrix exponential of the inpur matrix at the same spatial location.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vExpm

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))
output = vExpm(inputs)

torchvectorized.vlinalg.vLogm(inputs: torch.Tensor, replace_nans=False)¶

Compute the matrix logarithm $\mathbf{M} = \mathbf{U} log(\mathbf{\Sigma}) \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices. replace_nans (bool) – If `True`, replace nans by 0
Returns:	Return a tensor with shape Bx9xDxHxW where every voxel is the matrix logarithm of the inpur matrix at the same spatial location.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vLogm

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))
output = vLogm(inputs)

torchvectorized.vlinalg.vTrace(inputs: torch.Tensor)¶

Compute the trace of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices.
Returns:	Return a tensor with shape Bx1xDxHxW where every voxel is the trace of the inpur matrix at the same spatial location.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vTrace

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))
output = vTrace(inputs)

torchvectorized.vlinalg.vDet(inputs: torch.Tensor)¶

Compute the determinant of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices.
Returns:	Return a tensor with shape Bx1xDxHxW where every voxel is the determinant of the inpur matrix at the same spatial location.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym
from torchvectorized.vlinalg import vDet

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))
output = vDet(inputs)

nn¶

class torchvectorized.nn.EigVals¶

Differentiable neural network layer (torch.nn.Module) that performs eigendecomposition on every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW and return the eigenvalues.

See Ionescu et al., Matrix backpropagation for deep networks with structured layers, CVPR 2015 for details on the gradients computation

forward(x: torch.Tensor)¶

Takes a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW and return a volume of their eigenvalues

Parameters:	x (torch.Tensor) – A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW
Returns:	A tensor with shape *(BDHW)x3** where every voxel’s channels are the eigenvalues of the inpur matrix at the same spatial location.
Return type:	torch.Tensor

class torchvectorized.nn.Expm¶

Differentiable neural network layer (torch.nn.Module) that performs matrix exponential on every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

See Ionescu et al., Matrix backpropagation for deep networks with structured layers, CVPR 2015 for details on the gradients computation

forward(x: torch.Tensor)¶

Compute the matrix exponential $\mathbf{M} = \mathbf{U} exp(\mathbf{\Sigma}) \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	x (torch.Tensor) – A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW
Returns:	A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.
Return type:	torch.Tensor

class torchvectorized.nn.Logm¶

Differentiable neural network layer (torch.nn.Module) that performs matrix logarithm on every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

See Ionescu et al., Matrix backpropagation for deep networks with structured layers, CVPR 2015 for details on the gradients computation

forward(x: torch.Tensor)¶

Compute the matrix exponential $\mathbf{M} = \mathbf{U} log(\mathbf{\Sigma}) \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	x (torch.Tensor) – A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW
Returns:	A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.
Return type:	torch.Tensor

class torchvectorized.nn.ExpmLogm¶

Differentiable neural network layer (torch.nn.Module) that performs consecutive matrix exponential and logarithm on every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

See Ionescu et al., Matrix backpropagation for deep networks with structured layers, CVPR 2015 for details on the gradients computation

forward(x: torch.Tensor)¶

Compute the matrix exponential $\mathbf{M} = \mathbf{U} exp(\mathbf{\Sigma}) \mathbf{U}^{\top}$ and the matrix logarithm $\mathbf{M} = \mathbf{U} log(\mathbf{\Sigma}) \mathbf{U}^{\top}$ of every voxel in a volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.

Parameters:	x (torch.Tensor) – A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW
Returns:	A volume of flattened 3x3 symmetric matrices of shape Bx9xDxHxW.
Return type:	torch.Tensor

utils¶

torchvectorized.utils.overload_diag(inputs: torch.Tensor)¶

Add an EPSILON to the diagonal of every 3x3 matrix represented by the 9 channels of an input of shape Bx9xDxHxW to improve numerical stability

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices.
Returns:	A volume of shape Bx9xDxHxW where each voxel represent a flattened 3x3 symmetric matrix.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym, overloadd_diag

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = overload_diag(sym(torch.rand(b, c, d, h, w)))

torchvectorized.utils.sym(inputs: torch.Tensor)¶

Symmetrizes every 3x3 matrix represented by the 9 channels of an input of shape Bx9xDxHxW by applying $\frac{1}{2}(\mathbf{X} + \mathbf{X}^{\top})$ .

Parameters:	inputs (torch.Tensor) – The input tensor of shape Bx9xDxHxW, where the 9 channels represent flattened 3x3 symmetric matrices.
Returns:	A volume of shape Bx9xDxHxW where each voxel represent a flattened 3x3 symmetric matrix.
Return type:	torch.Tensor

Example:

import torch
from torchvectorized.utils import sym

b, c, d, h, w = 1, 9, 32, 32, 32
inputs = sym(torch.rand(b, c, d, h, w))

What is Torch Vectorized ?¶

vSymEig¶

vExpm¶

vLogm¶

Install me¶

How to use¶

vlinalg¶

nn¶

utils¶

Indices and tables¶