RuntimeError when using multiple GPUs

[yongtaowang327]

Hi, thanks for releasing the code!

When I tried to use multiple GPUs by model = RotoGrad(backbone, heads, 64, normalize_losses=True), model = torch.nn.DataParallel(model, device_ids=[0,1]).to(torch.device("cuda:0")), there was a RuntimeError:

  File "/data2/wyt/model/monodle_mtl/lib/helpers/trainer_helper.py", line 257, in train_one_epoch
    outputs = self.model(inputs)
  File "/data2/wyt/anaconda3/envs/monodle_mtl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data2/wyt/anaconda3/envs/monodle_mtl/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 168, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/data2/wyt/anaconda3/envs/monodle_mtl/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 178, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/data2/wyt/anaconda3/envs/monodle_mtl/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/data2/wyt/anaconda3/envs/monodle_mtl/lib/python3.7/site-packages/torch/_utils.py", line 461, in reraise
    raise exception
RuntimeError: Caught RuntimeError in replica 0 on device 0.

And when I tried to use single GPU by model = model.to(torch.device("cuda:0")), it looks like no problem.

I'm not sure if this project can support multiple GPUs, or if there is something wrong with my own code. Any help woud be greatly appreciated. Thanks!

[adrianjav]

Hi @yongtaowang327,
Sadly, the short answer is no, RotoGrad does not support multiple GPUs.

Personally, I have never needed to use multiple GPUs on my work, so I never thought of adding such a feature to the library. If you can point me out to the requirements to implement such a feature and the changes required are not too complicated (or, even better, if you feel like creating a PR that fixes it), I'd be happy to look into adding support to multiple GPUs.