Modify existing Pytorch code to run on multiple GPUs
I'm trying to run Pytoch UNet from the following link on 2 or more GPUs
Pytorch-UNet github
the changes the I did till now is:
1.
from:
net = UNet(n_channels=3, n_classes=1, bilinear=True)
logging.info(f'Network:\n'
             f'\t{net.module.n_channels} input channels\n'
             f'\t{net.module.n_classes} output channels (classes)\n'
             f'\t{&quot;Bilinear&quot; if net.module.bilinear else &quot;Transposed conv&quot;} upscaling')

to:
net = UNet(n_channels=3, n_classes=1, bilinear=True)
net = nn.DataParallel(net)
logging.info(f'Network:\n'
             f'\t{net.module.n_channels} input channels\n'
             f'\t{net.module.n_classes} output channels (classes)\n'
             f'\t{&quot;Bilinear&quot; if net.module.bilinear else &quot;Transposed conv&quot;} upscaling')




in each place where was:
net.&lt;something&gt;

replaced to:
net.module.&lt;something&gt;

I know that pytorch see more that 1 GPU because torch.cuda.device_count() return

2

.
But as long a I try to run train that need more momery than what the first GPU have I'm getting:

RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0;
11.91 GiB total capacity; 10.51 GiB already allocated; 82.56 MiB free; 818.92 MiB cached)

I change required memory to train by changing the batch size.
Any help welcome
EDIT
I see that training run twice faster with 2 GPUs, but max batch size of run with single GPU is the same as for two GPU. Is there any way to use the memory of 2 GPUs together during a single training run?

My mistake was changing output = net(input) (commonly named as model) to:
output = net.module(input)
you can find information here
