ahendriksen / tomosipo

A library for painless 3D tomographic reconstruction
https://aahendriksen.gitlab.io/tomosipo/index.html
GNU General Public License v3.0
83 stars 16 forks source link

Multi-GPU support ? #34

Open svvenkatakrishnan opened 2 days ago

svvenkatakrishnan commented 2 days ago

Does tomosipo support multiple-GPU usage like the ASTRA toolbox ?

askorikov commented 2 days ago

Hi! Yes, since ASTRA is the backend for tomosipo, you configure it in the same way. See here: https://aahendriksen.gitlab.io/tomosipo/topics/operator.html#large-datasets-and-multiple-gpus

svvenkatakrishnan commented 1 day ago

Thank you!

Is there an example script of integrating this with PyTorch i.e where we process each chunk on different GPUs ?

To put this another way: An example of processing a large volume - say reconstruction followed by a GPU based denoising ?

askorikov commented 1 hour ago

How do you implement multi-GPU processing on the PyTorch part? Integration with Astra will depend on that I think.

svvenkatakrishnan commented 28 minutes ago

For now, I dont have a script. More of a hypothetical with a goal to use this as a Plug-and-Play prior for large 3D volumes while avoiding data-movement between CPU and GPU.

So if we have some ||y-Ax||_{2}^2 + ||x-z||_2^{2} that we can implement with tomoSIPO where there was infrastructure to process different chunks on different GPUs, then in the next step we could call the denoisers to run independently on separate GPUs.