Neural network approximators of linear algebra operations on GPU with PyTorch
View Example notebook
The neuralg
module is a neural-network-based collection of approximators for common numerical linear algebra operations such as eigenvalue or singular value computation. It allows utilizing GPUs for differentiable and efficient computation with PyTorch. This also supports maintaining stable gradient flows.
The software is free to use and is designed for the machine learning community in general, and users focusing on topics involving numerical linear algebra in particular.
This project is built with the following packages:
Below is a short guide for setting up neuralg on your device.
Users who want to use the GPU needs to set up CUDA including installing a cudatoolkit.
To install neuralg via pip, run
pip install neuralg
Note that pip will not set up PyTorch with CUDA and GPU support.
GPU Utilization To set up the GPU version of PyTorch, please refer to installation procedures at PyTorch Documentation
The neuralg module is designed to resemble existing, commonly used numerical linear algebra libraries. Below is a small example showing how neuralg can be used to approximate the eigenvalues of a batch of random matrices. For a more elaborate and interactive example, please refer to the jupyter notebook example problem.
import torch
import neuralg
from neuralg import eigvals
# Enable GPU support if available
neuralg.set_up_torch(torch_enable_cuda = True)
# Sample a batch of matrices with uniform iid coefficients
# Note that neuralg only supports input of tensor type
batch_size, matrix_size = 10000, 5
matrix_batch = torch.rand(batch_size,matrix_size,matrix_size)
# Call neuralg to approximate eigenvalues
eigvals = eigvals(matrix_batch) # outputs are torch.complexfloat dtype
# Construct symmetric matrices
symmetric_matrix_batch = torch.triu(matrix_batch, 0) + torch.transpose(torch.triu(matrix_batch, 1),1,2)
# Symmetric matrices have specialized appproximators
eigvals_sym = eigvals(symmetric_matrix_batch, symmetric = True) # output are torch.float dtype
Built with PyTorch and targeting GPU utilization, neuralg only supports input of type torch.Tensor . The current version of neuralg supports real valued input matrices. Current available models have been trained and evaluated on random square matrices. For details, see specifications for corresponding operations and matrix types. |
function call | supported dtypes | supported sizes | training matrix distribution |
---|---|---|---|---|
eig(x,symmetric = True) |
float32 (in / out), float64 (in / out) | [2,20] | Real valued symmetric with i.i.d. centered normally distributed eigenvalues with variance 100/3 and eigenvectors uniformly distributed on the unit sphere. | |
eig(x,real = True) |
float32 (in / out), float64 (in / out) | [2,10] | Real valued asymmetric with i.i.d. centered normally distributed eigenvalues with variance 100/3. | |
eig(x) |
float32 (in), float64 (in), complexfloat64 (out), complexfloat128(out) | [2,5] | Real valued with i.i.d. uniformly distributed elements on [-10,10]. | |
svd(x) |
float32 (in / out), float64 (in / out) | [2,20] | Real valued with i.i.d. uniformly distributed elements on [-10,10]. |
All current available models are listed below. The nerf type models are implementations of architecture from NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis | model name | model type | hyper parameters | Additional information |
---|---|---|---|---|
EigNERF | nerf | 8 hidden layers, 200 neurons, skips at layers [2,4,6] | ||
CEigNERF | complex nerf, based on nerf | 8 hidden layers, 200 neurons, skips at layers [2,4,6] | Outputs are complex valued |
Additionally, neuralg supports training models from scratch or re-training and fine tuning existing models, depending on specific user applications. It is also possible to introduce new model architectures to expand the module. Please refer to the tutorial for a thorough how-to guide.
The project is open to community contributions. Feel free to open an issue or write us an email if you would like to discuss a problem or idea first.
If you want to contribute, please proceed according to:
git clone https://github.com/esa/neuralg.git
neuralg
conda env create -f environment.yml
conda activate neuralg
After cloning the repository, developers can check the functionality of neuralg
by running the following command in the root directory:
pytest
Now you are all set to contribute. PRs should be created from and into the develop
branch. For each release the develop branch is merged into main.
git checkout -b feature/NewFeature
)git commit -m 'Add some NewFeature'
)git push origin feature/NewFeature
)develop
branch, not main
and we will have a look at your contribution as soon as we can.
Furthermore, please make sure that your PR passes all automated tests, after which the contribution will be revire. Only PRs created on the develop
branch with all tests passing will be considered.
Distributed under the GPL-3.0 License. See LICENSE for more information.