Multi-GPU support ?
Closed this issue · 5 comments
Does tomosipo support multiple-GPU usage like the ASTRA toolbox ?
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
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 ?
How do you implement multi-GPU processing on the PyTorch part? Integration with Astra will depend on that I think.
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.
Well, most likely you will have to launch different denoisers from different processes, in which case Astra + tomosipo should work out of the box. You just need to set astra.set_gpu_index() in each process to the desired GPU.