Tensor-Train decomposition in pytorch
Tensor-Train decomposition package written in Python on top of pytorch
. Supports GPU acceleration and automatic differentiation.
It also contains routines for solving linear systems in the TT format and performing adaptive cross approximation (the AMEN solver/cross interpolation is inspired form the MATLAB TT-Toolbox).
Some routines are implemented in C++ for an increased execution speed.
Following requirements are needed:
python>=3.6
torch>=1.7.0
numpy>=1.18
opt_einsum
The GPU (if available) version of pytorch is recommended to be installed. Read the official installation guide for further info.
You can install the package using the pip
command:
pip install git+https://github.com/ion-g-ion/torchTT
One can also clone the repository and manually install the package:
git clone https://github.com/ion-g-ion/torchTT
cd torchTT
python setup.py install
TODO
The main modules/submodules that can be accessed after importing torchtt
are briefly desctibed in the following table.
Detailed description can be found here.
Component | Description |
---|---|
torchtt |
Basic TT class and basic linear algebra functions. |
torchtt.solvers |
Implementation of the AMEN solver. |
torchtt.grad |
Wrapper for automatic differentiation. |
torchtt.manifold |
Riemannian gradient and projection onto manifolds of tensors with fixed TT rank. |
torchtt.nn |
Basic TT neural network layer. |
torchtt.interpolate |
Cross approximation routines. |
The directory tests/ from the root folder contains all the unittests
. To run them use the command:
pytest tests/
The documentation ca be found here. Following example scripts (as well as python notebooks) are also provied provided as part of the documentation:
torchtt
(Try on Google Colab). mnist_nn.py / mnist_nn.ipynb: Example of TT layers used for image classification (Try on Google Colab).
The documentation is generated using shpinx
with:
make html
after installing the packages
pip install sphinx sphinx_rtd_theme
Ion Gabriel Ion, e-mail: ion.ion.gabriel@gmail.com