wasserth / TotalSegmentator

Tool for robust segmentation of >100 important anatomical structures in CT and MR images
Apache License 2.0
1.45k stars 245 forks source link

Add Python API support for multiple GPUs #369

Closed QianMuXiao closed 1 week ago

QianMuXiao commented 1 week ago

I often need to use TotalSegmentator to perform preliminary segmentation on large datasets to create rough labels, which is a preprocessing method I frequently use in my current research. However, processing hundreds of patients one by one using TotalSegmentator on a single GPU is very slow and often takes several days. My school provides multi-GPU servers, so I’m wondering if the Python API can support multi-GPU acceleration?

wasserth commented 1 week ago

The best way to utilise multiple GPUs for TotalSegmentator is to run it for several subjects in parallel on different GPUs. For TotalSegmentator itself it does not make sense to support MultiGPU since it anyways uses the GPU only for a few seconds.