Using dataparallel for custom classes

s_n: <p>Hi,<br>
I have a class for my model that uses some other model in it. For example:</p>
<p>class Mynet1(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(NN_model, self).<strong>init</strong>()<br>
self.fc1 = nn.Linear(4160, 500)</p>
<p>class Mynet2(nn.Module):<br>
def <strong>init</strong>(self, Mynet):<br>
super(NN_model, self).<strong>init</strong>()<br>
self.fc1 = nn.Linear(4160, 500)<br>
self.layer = Mynet(…)</p>
<p>mynet1=Mynet1()<br>
model = torch.nn.Dataparallel(Mynet2(mynet1))<br>
model.cuda()<br>
So if I want to run Mynet2 on multiple devices. I will do torch.nn.Dataparallel(Mynet2). Do I also need to take Mynet1 to multiple devices as in do torch.nn.Dataparallel(Mynet1) and then send it to Mynet2?<br>
Or it will automatically be on multiple devices since I take Mynet2 to multiple devices.</p>
<p>One more question related to dataparallel. Let’s say I have my custom conv2d class that uses torch variables like torch.mul, torch.sub etc. So when I was running this code on a single device I used to do “.to(device)” for all the intermediate variable in my custom conv2d class. So now when I run a model using this custom conv2d class on multiple devices using dataparallel I keep on getting this error: <strong>“RuntimeError: Expected tensor for argument <span class="hashtag">#1</span> ‘input’ to have the same device as tensor for argument <span class="hashtag">#2</span> ‘weight’; but device 0 does not equal 1 (while checking arguments for cudnn_batch_norm)”.</strong></p>
<p>I have taken my model that uses this custom conv2d class to multiple devices using torch.nn.dataparallel. So I don’t understand why I keep on getting this error. Is this because the intermediate variable that I have inside my custom conv2d class are on single device?<br>
If so how can I take this intermediate variable to multiple devices?<br>
Should I define them as model parameters using torch.nn.Parameter?</p>
<p>Thanks</p>

ptrblck: <aside class="quote no-group" data-username="s_n" data-post="1" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>So if I want to run Mynet2 on multiple devices. I will do torch.nn.Dataparallel(Mynet2). Do I also need to take Mynet1 to multiple devices as in do torch.nn.Dataparallel(Mynet1) and then send it to Mynet2?<br>
Or it will automatically be on multiple devices since I take Mynet2 to multiple devices.</p>
</blockquote>
</aside>
<p><code>nn.DataParallel</code> will automatically create the model replica with all internal submodules, parameters, and buffers on the specified devices, so you can wrap the main model only into <code>nn.DataParallel</code>.</p>
<aside class="quote no-group" data-username="s_n" data-post="1" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>Is this because the intermediate variable that I have inside my custom conv2d class are on single device?<br>
If so how can I take this intermediate variable to multiple devices?<br>
Should I define them as model parameters using torch.nn.Parameter?</p>
</blockquote>
</aside>
<p>It depends, how you’ve created these tensors.<br>
All registered <code>nn.Parameters</code> and buffers (via <code>self.register_buffer('name', torch.randn(1))</code>) will be pushed automatically to the corresponding device.<br>
However, if you have created a plain tensor and pushed it to a device via <code>to('cuda')</code>, you would either have to use the <code>device</code> argument of another parameter or register it as an <code>nn.Parameter</code> (if it needs gradients) of buffer.</p>

s_n: <p>Thanks for the reply <a class="mention" href="/u/ptrblck">@ptrblck</a>. So I won’t even have to take mynet1 instance to the device before passing it as an argument to Mynet2?</p>

s_n: <p>Regarding the answer to my second question. I actually don’t need the gradients for these intermediate variables. Can you please explain a bit more on what you mean by " However, if you have created a plain tensor and pushed it to a device via  <code>to('cuda')</code> , <strong>you would either have to use the  <code>device</code>  argument of another parameter</strong> or register it as an  <code>nn.Parameter</code>  (if it needs gradients) of buffer"</p>
<p>One solution to this problem I feel would be to define all my intermediate variables as nn.Parameters and make require_grad=False. Also, I can define them self.register_buffer and make grad=False. Will both solutions be similar?</p>

ptrblck: <aside class="quote no-group" data-username="s_n" data-post="3" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>So I won’t even have to take mynet1 instance to the device before passing it as an argument to Mynet2?</p>
</blockquote>
</aside>
<p>That is correct. Your <code>mynet1</code> would act as a “standard” module, e.g. similar to <code>nn.Linear</code> or <code>nn.Conv2d</code>, so you can just pass it into the parent model (<code>Mynet2</code>) and call <code>to()</code> or <code>cuda()</code> on the parent.</p>
<aside class="quote no-group" data-username="s_n" data-post="4" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>Can you please explain a bit more on what you mean by</p>
</blockquote>
</aside>
<p>All <code>nn.Modules</code>, <code>nn.Parameters</code> and buffers will be properly moved to the device, if <code>to()</code> or <code>cuda()/cpu()</code> is called on the parent device, while plain tensors will not. This code snippet demonstrates the behavior:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.module = nn.Linear(1, 1)
        self.param = nn.Parameter(torch.randn(1, 1))
        self.register_buffer('buf', torch.randn(1, 1))
        self.ten = torch.randn(1, 1)

    def forward(self, x):
        return self.module(x) + self.param + self.buf + self.ten
        
model = MyModel()
print(model.module.weight.device)
print(model.param.device)
print(model.buf.device)
print(model.ten.device)
out = model(torch.randn(1, 1)) # works
&gt; cpu
cpu
cpu
cpu

model.to('cuda')
print(model.module.weight.device)
print(model.param.device)
print(model.buf.device)
print(model.ten.device)
out = model(torch.randn(1, 1).to('cuda')) # fails, as self.ten is still on the CPU
&gt; cuda:0
cuda:0
cuda:0
cpu
RuntimeError: expected device cuda:0 but got device cpu
</code></pre>
<aside class="quote no-group" data-username="s_n" data-post="4" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>One solution to this problem I feel would be to define all my intermediate variables as nn.Parameters and make require_grad=False. Also, I can define them self.register_buffer and make grad=False. Will both solutions be similar?</p>
</blockquote>
</aside>
<p>If you don’t want to train this tensor, use a buffer. Using a parameter with <code>requires_grad=False</code> would also work, but a buffer would be the cleaner approach.</p>

s_n: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Actually for my conv2d function I am using autograd Functions. Like below</p>
<p>class Conv2d_function(Function):</p>
<pre><code>@staticmethod
def forward(ctx, x1,x2):
    
    y1 = torch.tensor([1,2,3]).to(device)
    y2 = torch.tensor([4,5,6]).to(device)
    ....
</code></pre>
<p>class Conv2d(nn.Module):<br>
def <strong>init</strong>(self, in_channels …):<br>
…<br>
…<br>
super(Conv2d, self).<strong>init</strong>()</p>
<pre><code>def forward(self, input):
    return Conv2d_function.apply(input)
</code></pre>
<p>Creating model using custom conv2d class<br>
class Model():<br>
layer1 = Conv2d(…)<br>
layer2 = Conv2d(…)</p>
<p>model1 = Model()</p>
<p>So when I use Dataparallel for this model1 I get this error “RuntimeError: Expected tensor for argument <span class="hashtag">#1</span> ‘input’ to have the same device as tensor for argument <span class="hashtag">#2</span> ‘weight’; but device 0 does not equal 1 (while checking arguments for cudnn_batch_norm)”</p>
<p>Is this error occurring because in the forward function of Conv2d_function class I am defining a tensor using .to(device)?</p>
<p>Actually the tensor y1 and y2 depend on my input to the forward function of class Conv2d so I can’t define those tensor in the <strong>init</strong> of Conv2d class as register_buffer or Parameter. So I can only define those in my forward function of Conv2d_function class. But with data-parallel I get this above error which I feel is because I am transferring these tensors to GPU in my forward pass.</p>

ptrblck: <aside class="quote no-group" data-username="s_n" data-post="6" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>Is this error occurring because in the forward function of Conv2d_function class I am defining a tensor using .to(device)?</p>
</blockquote>
</aside>
<p>Yes, and since you are not specifying the device id, the default device will be used in each model replica.</p>
<p>You could use the <code>device</code> attribute of the incoming tensors instead:</p>
<pre data-code-wrap="python"><code class="lang-python">y1 = torch.tensor([1,2,3]).to(x1.device)
y2 = torch.tensor([4,5,6]).to(x1.device)
</code></pre>

s_n: <p>Ok. A quick question in line with this. So when we pass input to our model. We run a for loop over the dataloader as shown below:<br>
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)<br>
for batch_idx,(data, target) in enumerate(testloader):<br>
data_var = data.to(device)<br>
target_var = target.to(device)<br>
output = model(data_var)<br>
So does dataparallel takes this input to all devices when we use data parallel?<br>
Because while passing input (data_var) to my model I take it to default device.<br>
Does pytorch create copies of input image as well on all the GPUs when we use dataparallel?</p>
<p>Because I want my tensors (y1, y2 …) to be on all the devices and now I will be initializing these tensors using x1.device</p>

ptrblck: <aside class="quote no-group" data-username="s_n" data-post="8" data-topic="83396">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/278dde/48.png" class="avatar"> s_n:</div>
<blockquote>
<p>So does dataparallel takes this input to all devices when we use data parallel?<br>
Because while passing input (data_var) to my model I take it to default device.<br>
Does pytorch create copies of input image as well on all the GPUs when we use dataparallel?</p>
</blockquote>
</aside>
<p><code>nn.DataParallel</code> will split the input tensor in dim0 and will send each chunk to the corresponding model replica on a specific device.<br>
I.e. your model’s forward will get an input of the shape <code>[batch_size//nb_gpus, *]</code>.<br>
If you need to create <code>y1</code> and <code>y2</code> inside the <code>forward</code> method, you should also consider the current batch size of the input.<br>
Also, it is correct to push the input tensor in the <code>DataLoader</code> loop to the default device, since <code>nn.DataParallel</code> will take care of the scattering and gathering of the tensors and parameters.</p>