SBU-BMI / wsinfer

🔥 🚀 Blazingly fast pipeline for patch-based classification in whole slide images
https://wsinfer.readthedocs.io
Apache License 2.0
56 stars 10 forks source link

[ENH] add ability to use multiple GPUs #102

Closed kaczmarj closed 1 year ago

kaczmarj commented 1 year ago

we might get speedups by using multiple gpus. this would require two changes:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

to use all available gpus. currently we specify "cuda:0" which uses the first available gpu.

second, we need to wrap the model in another class:

model = nn.DataParallel(model)

a good test of this would be the TIL model because it will have the most patches per slide (because the patch sizes are small).