Sending a tensor to multiple GPUs

semihcanturk: <p>I have a DataParallel model with a tensor attribute I need to define <em>after</em> I wrap the model with DataParallel. While the model has cuda device_ids = [0, 1] as expected, the tensor I assign to the model has device cuda:0 only, so it is not copied to all devices when I send it to the model. Is it possible to have this tensor available in both devices?</p>

Deepali: <p>To copy the tensor to the GPU you can use - <code>data = data.to(device)</code></p>
<p>Details at - <a href="https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html</a></p>

semihcanturk: <p>That would send the data from one GPU to another, though. I want the tensor to be on both GPUs, cuda:0 <em>and</em> cuda:1, just like the rest of the DataParallel model.</p>

Andreas_Georgiou: <p>Where do you create this tensor? Inside the wrapped model’s forward? Then assuming that <code>x</code> is the input to the model then you should use <code>x.new_tensor()</code> or <code>x.new_zeros()</code> or <code>torch.zeros_like()</code> to create the tensor automatically in the correct device.</p>
<p>See here: <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.new_zeros" rel="nofollow noopener">https://pytorch.org/docs/stable/tensors.html#torch.Tensor.new_zeros</a></p>
<p>If you are assigning the attribute from the outside e.g <code>model.tensor = some_tensor</code> then this is not going to work. The model is replicated to each device after each iteration. Doing <code>model.tensor</code> does not assign the attribute to the real model inside the wrapper but directly to the <code>DataParallel</code> instance.</p>

semihcanturk: <p>What I am doing is in every batch I assign the tensor (named <code>graph</code>) to the module wrapped in DataParallel based on a criterion. This tensor is different for different batches and should not be parallelized, which is why I am trying to assign it directly to the model:</p>
<pre><code class="lang-auto">        for batch_idx, (x, y, graph, subject) in enumerate(self.train_loader):

            if model.module.subject != subject:

                model.module.subject = subject
                model.module.graph = graph      # where the assignment takes place

            output = model(x.to(self.device))
            target = torch.argmax(y, dim=1)

            optimizer.zero_grad()
            loss = F.nll_loss(output, target, weight=self.w)
            loss.backward()
            optimizer.step()
</code></pre>
<p>Given that I assign the tensor to the <em>module</em> wrapped by DataParallel, I would expect this tensor to be replicated in each GPU as well, which does not seem to be the case. This might have to do with the replication process:</p>
<aside class="quote no-group" data-username="Andreas_Georgiou" data-post="4" data-topic="49390">
<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/andreas_georgiou/48/11871_2.png" class="avatar"> Andreas_Georgiou:</div>
<blockquote>
<p>The model is replicated to each device after each iteration.</p>
</blockquote>
</aside>
<p>When exactly does the replication happen? And if this happens in every iteration, shouldn’t the model replicate this tensor as well?</p>

Andreas_Georgiou: <p>Try using <code>self.register_buffer('graph', None)</code> inside <code>__init__</code> of the model. This way <code>DataParallel</code> knows that this is a tensor that must be copied too. <code>DataParallel</code> only replicates parameters and buffers.</p>

semihcanturk: <p>I think that solves it, thank you!</p>
<p>One issue is even though the outputs and the subsequent loss is now calculated, the code now gets stuck on the backward pass. Not entirely sure whether this issue is directly related, though I do not encounter it in the single-GPU case.</p>

Andreas_Georgiou: <p>That could be an issue with your setup. Maybe hardware. I don’t think it is related to pytorch</p>