Using DataParallel when the input to the model is a dict

Jimmy2027: <p>Hello,<br>
I am trying to train a multi-modal model on multiple GPUs using torch.nn.DataParallel.<br>
However, I have multiple modalities so the input to the model is a dictionary.<br>
Is there any way to make this work on multiple GPUs? As far as I’ve understood DataParallel only works if the input to the model is a tensor.</p>
<p>My input to the model looks like the following:</p>
<pre><code class="lang-auto">{
'modality_1': torch.Tensor((bs, *img_size)),
'modality_2':  torch.Tensor((bs, *img_size)),
'modality_3':  torch.Tensor((bs, *img_size)),
}
</code></pre>

klory: <p>Sol 1: define your own class for your inputs and inplement the <code>to()</code> function</p>
<p>Sol2:</p>
<pre><code class="lang-auto">self.data = {k: v.to(device) for k, v in self.data.items()}
</code></pre>
<p>Ref:<br>
</p><aside class="onebox allowlistedgeneric">
  <header class="source">
      <img src="https://huggingface.co/transformers/_modules/transformers/tokenization_utils_base.html/../../_static/favicon.ico" class="site-icon" width="" height="">
      <a href="https://huggingface.co/transformers/_modules/transformers/tokenization_utils_base.html" target="_blank" rel="noopener nofollow ugc">huggingface.co</a>
  </header>
  <article class="onebox-body">
    <img src="" class="thumbnail" width="16" height="16">

<h3><a href="https://huggingface.co/transformers/_modules/transformers/tokenization_utils_base.html" target="_blank" rel="noopener nofollow ugc">transformers.tokenization_utils_base — transformers 3.5.0 documentation</a></h3>



  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Jimmy2027: <p>Hi <a class="mention" href="/u/klory">@klory</a>, thanks for your answer!<br>
I’m not quite sure if I understand correctly: should I send each value in the dict to another GPU for parallel computing?</p>

klory: <p>You only need to call <code>nn.DataParallel</code> to your model, the data will be automatically distributed to all GPUs as long as you called <code>to()</code></p>

Jimmy2027: <p>Oh okay I see, looks easy <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
I’m getting following error though on the line <code>results = model(batch)</code>:</p>
<pre><code class="lang-auto">  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 156, in forward
    return self.gather(outputs, self.output_device)
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 168, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 68, in gather
    res = gather_map(outputs)
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in gather_map
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in &lt;genexpr&gt;
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in gather_map
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 61, in &lt;genexpr&gt;
    return type(out)(((k, gather_map([d[k] for d in outputs]))
  File "miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 63, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
TypeError: 'Laplace' object is not iterable
</code></pre>
<p>Is there anything I could be doing wrong?<br>
I have applied following line to the batch:<br>
<code>   batch = {k: v.to(torch.device('cuda')) for k, v in batch.items()}</code></p>

klory: <aside class="quote no-group" data-username="Jimmy2027" data-post="5" data-topic="102281">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jimmy2027/48/50502_2.png" class="avatar"> Jimmy2027:</div>
<blockquote>
<p><code>parallel</code></p>
</blockquote>
</aside>
<p>You did not show the line that triggers the mistake, can you run it successfully on one GPU? or (CPU at least)?</p>

Jimmy2027: <p>When trying to run this with <code>torch.device('cpu')</code> I get the error:</p>
<pre><code class="lang-auto">  File "/src/MIMIC/mimic/run_epochs.py", line 94, in basic_routine_epoch
    results = model(batch)
  File "/miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/miniconda3/envs/mimic/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 147, in forward
    raise RuntimeError("module must have its parameters and buffers "
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu
</code></pre>

klory: <p>how about GPU without nn.DataParallel?</p>

Jimmy2027: <p><a class="mention" href="/u/klory">@klory</a> the code runs without problems  on gpu without nn.DataParallel. The problem seems to be the <code>torch.distributions.Laplace</code> that I call in the forward pass. Is there any reason this should be a problem?</p>
<p>Also my model consists of different nn.Modules: an encoder and a decoder for each modality. Is this compatible with nn.DataParallel?</p>

klory: <p>Don’t know too much about Laplace,</p>
<p>But two modules are OK</p>

Jimmy2027: <p>I have created a <a href="https://gist.github.com/Jimmy2027/a8b460180bc8bca757da4cc647df6595" rel="noopener nofollow ugc">gist</a> that reproduces the error. The problem is indeed the Laplace distribution from <code>torch.distributions</code></p>