Multi-GPU Training

[michaelneumeier]

Multi-GPU Training

[michaelneumeier]

Thank you @bamsumit for the quick reply.

I cannot run the exact code that you proposed. For this I get a device mismatch between the input tensor and weight tensor. Thus I tried this:

net = Network().to(device)
net.forward(torch.rand(1, 2, 128, 128, 1).to(device))
net = nn.DataParallel(net, device_ids=[0, 1])

Then this code snippet works.

If I run it before the training loop, I can even run the training for several epochs without receiving the backprop error message from the previous comment.
But without the net.forward line, the error is still there. Do you understand why this is the case?

So thank you for the hint, I can now train on multiple GPUs.

[bamsumit]

Hi @michaelneumeier can you try this.

net = Network()
net.forward(torch.rand(1, C, H, W, 1).to(device))  # CHW = your input dimension
net = nn.DataParallel(net, device_ids=[0, 1])
net.to(device)

[bamsumit]

Glad you can run it now. The slayer models support runtime shape identification. This means that we don't need to specify the shape of a layer. Instead they are identified on runtime. With multi GPU, the shape update happens on the spawned dataparallel copy but is not reflected back on the main model, so it needs to be done on the main copy once with a dummy load before the multiGPU run. This should also work.

net = Network().to(device)
net = nn.DataParallel(net, device_ids=[0, 1])
...
net.module.forward(torch.rand(1, 2, 128, 128, 1).to(device))
# training sequence