JuliaWolleb / Diffusion-based-Segmentation

This is the official Pytorch implementation of the paper "Diffusion Models for Implicit Image Segmentation Ensembles".
MIT License
271 stars 35 forks source link

Run on Multiple GPUS #43

Closed TheRealCameronEvans closed 8 months ago

TheRealCameronEvans commented 10 months ago

Hello, thanks for sharing your code! Would it be possible for me to run the processes across multiple GPUs? I have a few low powered ones so this would speed my jobs up a lot. Thanks!

JuliaWolleb commented 10 months ago

Good morning

Yes, to run across multiple GPUs, this is already implemented in the code. You will need to set the flag use_ddp = True, you will enter the torch DistributedDataParallel (DDP) in the file guided_diffusion/train_util.py .

Let me know whether this answers your question. Best, Julia