Write operators only once and use it everywhere.
Write an operator in C++/CUDA and generate wrappers to different languages such as Python and machine learning libraries such as Tensorflow or Pytorch.
optox provides a tensor interface to ease data transfer between host tensors optox::HTensor and device tensors optox::DTensor of any floating type and number of dimensions.
Using this interface, an operator is only written once in C++/CUDA and wrappers for Python, Tensorflow 2.x and Pytorch expose the functionality to a higher level application (e.g. iterative reconstruction, custom deep learning reconstruction, ...).
The source files are organized as follows:
.
+-- src : `optox` library source files
| +-- tensor : header only implementation of `HTensor` and `DTensor`
| +-- operators : actual implementation of operator functionality
+-- python : python wrappers
+-- pytorch : pytorch wrappers
+-- tensorflow : tensorflow wrappers
We provide a python installation script. To build the respective packages for python, pytorch and tensorflow, simply add the flags --python, --pytorch, --tensorflow.
To build optox with gpuNUFFT, you only need to add the flag --gpunufft.
git clone https://github.com/midas-tum/optox.git
cd optox
# build all
python install.py --python --pytorch --tensorflow --gpunufft
# build only with pytorch support
python install.py --pytorch
We highly recommend to use Cuda >11.1 which runs smoothly with both Pytorch and Tensorflow
First setup the following environment variables:
CUDA_BIN_PATHto point to the NVidia CUDA toolkit (typically/usr/local/cuda)
Note that the CUDA version used to build the optox library should match the version required by Tensorflow and/or Pytorch.
We provide an anaconda environment with Tensorflow 2.4, Pytorch 1.9, Cuda 11.1. The environment optox can be created via
conda env create -f environment.yml
To build the basic optox library perform the following steps:
mkdir build
cd build
cmake ..
make installIf the default gcc compiler version is >8, you will get a build error. A simple workaround is to call following before the cmake command and with a clean build dir, assuming gcc-8 and g++-8 are installed on your system in /usr/bin/
export CC=/usr/bin/gcc-8
export CXX=/usr/bin/g++-8
If you experience a problem that lcudart library cannot be found, then link build to the location of the CUDA lib64 libraries (typically <path-to-cuda-lib64> = /usr/local/cuda/lib64)
export LDFLAGS=-L<path-to-cuda-lib64>
Use CMAKE_BUILD_TYPE=Release to avoid the device synchronization after each CUDA call. Then, no CUDA errors are generated but runtime is strongly reduced.
To build the Python wrappers optox requires pybind11 which can be installed in an anaconda environment by conda install pybind11.
To also build Python wrappers substitute the cmake command by:
cmake .. -DWITH_PYTHON=ONTo build it, the pytorch package must be installed.
cmake .. -DWITH_PYTORCH=ONTo build it, the tensorflow package must be installed.
cmake .. -DWITH_TENSORFLOW=ONNote that multiple combinations are supported.
optox provides python, pytorch and tensorflow/keras wrappers for khammernik/gpuNUFFT branch cuda_streams. Make sure to build gpuNUFFT in Release mode.
git clone --branch cuda_streams https://github.com/khammernik/gpuNUFFT.git
cd gpuNUFFT/CUDA
mkdir -p build
cd build
cmake .. -DGEN_MEX_FILES=OFF -DCMAKE_BUILD_TYPE=Release
make
To build them with optox, the gpuNUFFT package must be built and the path GPUNUFFT_ROOT_DIR (e.g. GPUNUFFT_ROOT_DIR=~/gpuNUFFT)
must be set.
cmake .. -DWITH_GPUNUFFT=ON
To perform an adjointness test of the nabla operator using the Python wrappers execute
python -m unittest optopy.nabla
If successful the output should be
(optox) ∂ python -m unittest optopy.nabla
dtype: <class 'numpy.float64'> dim: 2 diff: 6.661338147750939e-16
.dtype: <class 'numpy.float64'> dim: 3 diff: 2.842170943040401e-14
.dtype: <class 'numpy.float32'> dim: 2 diff: 2.86102294921875e-06
.dtype: <class 'numpy.float32'> dim: 3 diff: 7.62939453125e-06
.
----------------------------------------------------------------------
Ran 4 tests in 1.099s
OK
To perform a gradient test of the activations operators using the Pytorch wrappers execute
python -m unittest optoth.activations.act
If successful the output should be
(optox) ∂ python -m unittest optoth.activations.act
grad_x: -3616.3090656 num_grad_x -3616.3090955 success: True
grad_w: 7232.6181312 num_grad_w 7232.6181312 success: True
.grad_x: 535.2185935 num_grad_x 535.2185935 success: True
grad_w: 2236.8791233 num_grad_w 2236.8791233 success: True
.grad_x: -215.0009414 num_grad_x -215.0009432 success: True
grad_w: 430.0018828 num_grad_w 430.0018828 success: True
.
----------------------------------------------------------------------
Ran 3 tests in 2.263s
OK
To perform an adjointness test of the nabla operators using the Tensorflow wrappers execute
python -m unittest optotf.nabla
If successful the output should be
(optox) ∂ python -m unittest optotf.nabla
...
dtype: <dtype: 'float64'> dim: 2 diff: 1.0658141036401503e-14
.dtype: <dtype: 'float32'> dim: 2 diff: 0.0
.
----------------------------------------------------------------------
Ran 2 tests in 1.490s
OK
The keras layers can be found in optotf.keras.xxx.
Unittests can be called as follows:
python -m unittest discover optotf.test
python -m unittest discover optoth.test
python -m unittest discover optopy.test