Why model.to(device) wouldn't put tensors on a custom layer to the same device?

yxchng: <p>Currently, I have to pass a <code>device</code> parameter into my custom layer and then manually put tensors onto the specified device manually using <code>.to(device)</code> or <code>device=device</code>.</p>
<p>Is this behavior expected? It looks kind of ugly to me.</p>
<p>Shouldn’t <code>model.to(device)</code> put all the layers, including my custom layer, to device for me?</p>

ptrblck: <p>Could you post a small code snippet reproducing this error?<br>
I tried to reproduce your error, but it seems to work fine:</p>
<pre><code class="lang-auto">class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc1 = nn.Linear(1, 1)
        
    def forward(self, x):
        x = self.fc1(x)
        return x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 1)
        self.module1 = MyModule()
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.module1(x)
        return x   

model = Net()
model = model.to('cuda:0')
print(model.module1.fc1.weight.type())
&gt; torch.cuda.FloatTensor
print(model.fc1.weight.type())
&gt; torch.cuda.FloatTensor
</code></pre>

yxchng: <pre><code class="lang-auto">import torch.nn as nn
import torch

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc1 = nn.Linear(1, 1)

    def forward(self, x):
        n = torch.range(0,5)
        print(n.type())
        #x = x.mm(n)
        x = self.fc1(x)
        return x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 1)
        self.module1 = MyModule()

    def forward(self, x):
        x = self.fc1(x)
        x = self.module1(x)
        return x

model = Net()
model = model.to('cuda:0')
model(torch.zeros(1, 10).to('cuda:0'))
print(model.module1.fc1.weight.type())
print(model.fc1.weight.type())
</code></pre>
<p>Output:</p>
<blockquote>
<p>torch.FloatTensor<br>
torch.cuda.FloatTensor<br>
torch.cuda.FloatTensor</p>
</blockquote>
<p>The new tensor I define in my custom layer is not on cuda. As a result, I need to pass in <code>device</code> so that I can move them to <code>device</code>.</p>
<p>I would expect the default to be that new tensor should also be on cuda once your model is moved to cuda.</p>

yxchng: <p>Ok this is not a good example because I could just move it to the same tensor as x.</p>
<p>But the problem I face is because I have to define the weights myself</p>
<p><code>self.w = Parameter(torch.zeros(out_features, in_features))</code></p>
<p>in __init__ function.</p>
<p>At that point, I still do not know the input device. Therefore my self.w is on the cpu but my input is on the gpu.</p>

SimonW: <p>parameters will be moved as well in <code>.to</code> so you don’t have to worry about that parameter you add.</p>

yxchng: <p>Ok I found out the problem. It is not with weight that has been declared as a Parameter. It is with a map that I define in __init__ function. I need this map to also be put into the same <code>cuda</code> device as weight but I do not want it to be a Parameter.</p>
<p>I tried having these lines</p>
<pre><code class="lang-auto">self.w = Parameter(torch.zeros(out_features, in_features))
self.k_map = torch.zeros(self.m+1, device=self.w.device)
</code></pre>
<p>Apparently at initialization, the weight is still on the <code>cpu</code>, which causes the map to be on the <code>cpu</code>.</p>
<p>How should I map them both on <code>cuda</code> but not making <code>k_map</code> a Parameter?</p>

justusschock: <p>You could derive the <code>.cuda()</code>  and/or the <code>.to()</code> function for something like this:</p>
<pre><code class="lang-auto">class Net(torch.nn.Module):
    def __init__(self):
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.zeros(1)

    def forward(self, inputs):
        #do some stuff
    
    def cuda(self, device=None):
        self = super().cuda(device)
        self.b = self.b.cuda(device)
        return self 

    def to(self, *args, **kwargs):
        self = super().to(*args, **kwargs) 
        self.b = self.b.to(*args, **kwargs) 
        return self

</code></pre>
<p>EDIT: I recommend only to derive the <code>.to()</code> function because you have to make sure that you also derive the <code>.cpu()</code> function when deriving <code>.cuda()</code></p>

ptrblck: <p>In that case you should use <a href="https://pytorch.org/docs/stable/nn.html?highlight=register_buffer#torch.nn.Module.register_buffer">register_buffer</a>.</p>

yxchng: <p>Thanks a lot! register_buffer is indeed what I am looking for.</p>

sudocoder: <p>Sorry if new to this, how does <code>register_buffer</code> handle the <code>.to(device)</code> problem here, such that the custom module goes on the same device as the model it is part of. Thank you so much!</p>

ptrblck: <p>The <code>.to()</code> method will be applied on all internal <code>_parameters</code> and <code>_buffers</code> as can be seen <a href="https://github.com/pytorch/pytorch/blob/cef0443464a4ff5e3fd2e3b6eca0ee76c5c428ce/torch/nn/modules/module.py#L432" rel="nofollow noopener">here</a> and <a href="https://github.com/pytorch/pytorch/blob/cef0443464a4ff5e3fd2e3b6eca0ee76c5c428ce/torch/nn/modules/module.py#L224" rel="nofollow noopener">here</a>.</p>
<p>The <code>_parameters</code> and <code>_buffers</code> attributes will be registered, if you set an <code>nn.Parameter</code> as an attribute or use <code>self.register_buffer</code>. Have a look at this example:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_param = nn.Parameter(torch.randn(1))
        self.register_buffer('my_buffer', torch.randn(1))
        self.neither_param_nor_buf = torch.randn(1)
        
        
model = MyModel()
print(dict(model.named_parameters()))
&gt; {'my_param': Parameter containing:
tensor([-1.1189], requires_grad=True)}
print(dict(model.named_buffers()))
&gt; {'my_buffer': tensor([-0.0459])}

model.to('cuda')
print(dict(model.named_parameters()))
&gt; {'my_param': Parameter containing:
tensor([-1.1189], device='cuda:0', requires_grad=True)}
print(dict(model.named_buffers()))
&gt; {'my_buffer': tensor([-0.0459], device='cuda:0')}
print(model.neither_param_nor_buf.device)
&gt; cpu
</code></pre>
<p>As you can see, only the parameters and buffers were moved, while the tensor is still on the CPU.</p>

Qing_En: <p>It works for me. Thank you very much. <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

jackeown: <p>Is there any way to make sure that tensors created in the forward method are also on the appropriate device without passing in the device explicitly? Do you need to register buffers in the forward pass in this case? Seems kinda weird…maybe you could say something like <code>myForwardTensor.to(self.device)</code>?</p>

ptrblck: <p>You could use the <code>device</code> attribute of the input tensor or of any parameter (<code>next(self.parameters()).device</code>).</p>

emmb: <p>solved my problem by seeing this! Thanks!</p>

jastern33: <p>One thing to be careful with here - make sure that the you don’t assign the tensor to a variable, and then register the variable as a buffer, like this:</p>
<pre><code class="lang-auto">#  Incorrect way
self.my_tensor = torch.randn(1)
self.register_buffer('my_buffer', my_tensor)
</code></pre>
<p>If you do this, then try to access the original variable instead of the buffer, you will find that the device has not propagated to the variable.</p>
<pre><code class="lang-auto">def forward(self, inputs):
    inputs = inputs / self.my_tensor # ERROR: self.my_tensor is on cpu
</code></pre>
<p>This is the right way to do it:</p>
<pre><code class="lang-auto">def __init__(self):
    self.register_buffer('my_buffer', torch.randn(1))
...
def forward(self, inputs):
    inputs = inputs / self.my_buffer # This works. self.my_buffer is on gpu
</code></pre>

ptrblck: <p>That’s a good point and thanks for sharing.<br>
Since you’ve tagged me I assume one of my code snippets shows this behavior? (I can’t find it here, so could you send me a link to it so that I could add a comment or correct it?)</p>

jastern33: <p>Nope, your code is correct. I was following your code and ran into this problem because I implemented register_buffer slightly differently than you did. So I thought I’d share to help others avoid my mistake <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>