banctilrobitaille / torch-vectorized Goto Github PK

Fast analytical implementation of batch eigen-decomposition for 3x3 symmetric matrices with Pytorch. > 250x faster than regular Pytorch implementation of batch eigen-decomposition on GPU.

Home Page: https://torch-vectorized.readthedocs.io/en/latest/

License: MIT License

Python 99.65% Batchfile 0.17% Shell 0.18%

eigen-decomposition pytorch analytical fast 3x3 matrices batch deep-learning volumetric-3x3-matrices batch-eigen-decomposition

torch-vectorized's Introduction

Torch Vectorized

Batched and vectorized operations on volume of 3x3 symmetric matrices with Pytorch. 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.

Read the documentaton HERE

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, eigenvectors=True)

Contributing

How to contribute ?

Create a branch by feature and/or bug fix
Get the code
Commit and push
Create a pull request

Branch naming

Feature branch

feature/ [Short feature description] [Issue number]

Bug branch

fix/ [Short fix description] [Issue number]

Commits syntax:

Adding code:

+ Added [Short Description] [Issue Number]

Deleting code:

- Deleted [Short Description] [Issue Number]

Modifying code:

* Changed [Short Description] [Issue Number]

Merging code:

Y Merged [Short Description] [Issue Number]

Icons made by Freepik from www.flaticon.com is licensed by CC 3.0 BY

torch-vectorized's People

Stargazers

Watchers

Forkers

cy69855522 zhenglong178 p0l0satik

torch-vectorized's Issues

Getting negative eigenvalues for PDS matrices

Always getting negative eigenvalues for randomly generated PDS matrices when one eigenvalue is very small and close to 0 or very big. Any way to deal with this situation? The following is one specific example that gives a negative eigenvalue and the eigenvalue is not really close to 0.

T = torch.tensor([[ 1.3999e+00,  1.5765e+00, -5.5419e+03],
        [ 1.5765e+00,  2.1994e+00, -7.3147e+03],
        [-5.5419e+03, -7.3147e+03,  2.4693e+07]], dtype=torch.float64) 
vlinalg.vSymEig(T.reshape((1,9,1,1,1)), eigenvectors=False)[0]

The result is
tensor([[[[[-7.9096e-02]]], [[[ 2.6781e-01]]], [[[ 2.4693e+07]]]]], dtype=torch.float64)

But if we use the build-in function in Pytorch:

torch.symeig(T, eigenvectors=True)[0]

Then the eigenvalues are all positive:
tensor([4.5855e-03, 1.8413e-01, 2.4693e+07], dtype=torch.float64)

back-propagation issue

Thanks for the code. I am trying to use it to accelerate eigen decomposition. My input matrices are in shape (B, 3 ,3). To apply vSymEig, I first reshape it into (B, 9), and then expand the dimension to (B, 9, 1, 1, 1), and pass it through vSymEig. Finally I unsqueeze the output to return back to (B, 3) and (B, 3, 3). The code is as follows:

A_flat = A.reshape(-1, 9)[..., None, None, None]
w, v = vSymEig(A_flat, eigen_vectors=True)
w = w.squeeze()
v = v.squeeze()

However when I run later loss.backward(), this piece of code gives me error as follows:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.DoubleTensor [10900, 3, 1, 1, 1]], which is output 0 of SqrtBackward, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

Do you happen to meet similar problem, or any suggestions?

RuntimeError: Function 'AcosBackward' returned nan values in its 0th output.

Getting this error during training. I use a symmetric matrix, overloaded the diagonal with epsilon as well, but still getting this error. Any idea what it might be from?

Possible to extend to 4x4 matrices?

Hi @banctilrobitaille , thank you for your open-source sharing and I found that this is really helpful to accelerate 3x3 batch eigendecomposition. Actually, I am facing a problem of calculating the eigenvalue and eigenvectors of a 61440x4x4 tensor, 61440 is batch size, 4x4 is matrix size. So I'm wondering whether this fast algorithm can be extend to 4x4 matrices? If so, could you provide me some detailed suggestions?

Not supporting 0s only 3x3 matrix

At the moment the algorithm does not support the decomposition of null matrices. The returned eigen-values and eigen-vectors for such matrices are NaNs.

The expected results are 0s eigen-values and identity matrix for eigen-vectors.