neuralg

Neural network approximators of linear algebra operations on GPU with PyTorch
View Example notebook

About

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.

Built With

This project is built with the following packages:

• PyTorch, implying full differentiability and can be used for machine learning

Getting Started

Below is a short guide for setting up neuralg on your device.

Prerequisites

Users who want to use the GPU needs to set up CUDA including installing a cudatoolkit.

Installation

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

Usage

Small example

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

Support and training distributions

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.

Available models

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

Customizing and adding models

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.

Goals

• Proof of concept: Recent years of rapid advances in machine learning aside, neural network models are yet to reach competitive results in the field of numerical linear algebra. Some attention has been paid to the subject, e.g. with parameter rich transformers. neuralg serves as a demonstration of a competitive small-scale approach, with the goal of mitigating issues with memory and time complexity related to larger models.
• Supporting science: Linear algebra problems serve as fundamental computational components in countless science and engineering applications. Eigenvalue and singular value decompositions, solving linear system of equations and matrix inversion appear as essential parts of solutions in optimization, dynamical systems, signal processing etc. Ultimately, neuralg aims to provide useful tools to researchers within these fields, with a focus on parallell computation.
• Addressing efficient vectorization: The nature of classical numerical algorithms for linear algebra operations are often iterative and difficult to vectorize efficiently on GPUs and specialized machine learning hardware. Existing built-in libraries often synchronize with CPU, which can severly slow down computation. To this end, neuralg aims to allow users to exploit the computational benefits from GPU parallelization on targeted hardware.

Contributing

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:

1. Fork the project on GitHub.
2. Get the most up-to-date code by installing neuralg from source:
   1. Get miniconda or similar
   2. Clone the repo

   git clone https://github.com/esa/neuralg.git
   

   3. Setup the environment. This will create a conda environment called 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.

3. Create your Feature Branch (git checkout -b feature/NewFeature)
4. Commit your Changes (git commit -m 'Add some NewFeature')
5. Push to the Branch (git push origin feature/NewFeature)
6. Open a Pull Request on the 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.

License

Distributed under the GPL-3.0 License. See LICENSE for more information.