Proper way to write modules portable with .to(device)

bbharland: <p>I would like to know the correct way to write a module that will avoid <code>RuntimeError: Expected all tensors to be on the same device</code>.  (I am using torch.jit.script if this is important)</p>
<p>I recently asked a related question about the .to(device) function and was helpfully told that “valid members for this operation is other nn.Modules, Parameters and Buffers”.  What should you do with other kinds of data?</p>
<p>Here is an illustration of what I would like to do:</p>
<pre><code class="lang-auto">class TestModule(nn.Module):

    def __init__(self, net, param_float, param_tensor):
        super().__init__()
        self.net = net
        self.tensor_ = nn.Parameter(param_tensor)
        self.float_ = nn.Parameter(torch.tensor(param_float))
        self.tuple_of_tuples_ = ((1, 4), (1, 5))
        self.tensor_list_ = []
        
    def forward(self, x):
        # something that uses all module attributes

</code></pre>
<p>Many thanks!</p>

ptrblck: <p>As already described, parameters, buffers, and other modules will be registered to the parent module and thus the <code>to()</code> calls will also be applied on these.<br>
Using plain tensors as attributes in your module (or <code>list</code>s etc.) won’t have such an effect.<br>
The main question would be why you would use these tensors, which are neither parameters nor buffers nor are they inputs to the <code>forward</code> method?</p>

bbharland: <p>In my forward function, the first step is a featurization which looks like this:</p>
<pre><code class="lang-auto">x = torch.stack([
       torch.sqrt(torch.sum(torch.square(r[j] - r[i])))
       for (i, j) in self.pairs
       ]).reshape(1, -1)
</code></pre>
<p>where <code>self.pairs</code> is a Tuple[Tuple[int]] (I’m not sure how to write that properly for type hints.)</p>
<p>Another example that has me stumped is a growing list of tensors.  Occasionally, during some forward steps, this list will have a tensor appended to it.  This list is used during calls to forward.</p>
<p>Lastly, I have constants.  For these, it appears I want <code>nn.Parameter(param, requires_grad=False)</code></p>
<p>Is it true that these kinds of features weren’t intended to be used in a nn.Module?  I might be missing something bigger.  Thanks for your help!</p>

ptrblck: <p>It depends a bit on your use case and where these tensors are coming from.<br>
Generally you could use:</p>
<ul>
<li>trainable parameters registering them to the model via <code>nn.Parameter(List)</code>
</li>
<li>Non-trainable tensors via <code>self.regiter_buffer</code> (or set the <code>requires_grad</code> attribute of a parameter to <code>False</code>): this would properly register the tensors into the model and they would also show up in the <code>state_dict</code>, so you could save/load them</li>
<li>Non-trainable tensors via <code>self.register_buffer(tensor, persistent=False)</code> which will be registered into the model, but won’t show up in the <code>state_dict</code>
</li>
<li>Pass tensors to the <code>forward</code> method (with the right <code>dtype</code> and on the right <code>device</code>
</li>
</ul>
<p>However, if you need to create new tensors in the <code>forward</code>, you could use the <code>.device</code> attribute of any parameter (assuming they are all pushed to the device already). In your use case I don’t know how <code>self.pairs</code> is treated, but as already described a plain <code>list</code> won’t be registered, so you would need to push these tensors to the device manually.</p>

bbharland: <p>Thank you again.  So for the case of the integer tuple, I can imagine a slightly hacky way to substitute a tensor with dtype int and register it as a buffer.</p>
<p>I’m still mystified about the growing list of tensors.  Could I <code>self.register_buffer</code> an empty tensor and safely reassign it during code execution during <code>forward</code> so that it can grow longer?  Like this</p>
<pre><code class="lang-auto">    self.register_buffer('a', a)

def forward(self, x: torch.Tensor, extend: bool):
    if extend:
        new_tensor = ...
        self.a = torch.stack(self.a, new_tensor)
</code></pre>

ptrblck: <p>I think this use case should be possible if you make sure the shapes are as expected.<br>
Here is a small example, which concatenates a new tensor to the internal buffer and checks that it’s also added to the <code>state_dict</code>:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('a', torch.ones(1, 1))

    def forward(self, x: torch.Tensor, extend: bool):
        if extend:
            new_tensor = torch.randn(1, 1)
            self.a = torch.cat([self.a, new_tensor], dim=0)
        return x

model = MyModel()
x = torch.randn(1, 1)

for _ in range(10):
    out = model(x, extend=True)
    print(model.state_dict())
</code></pre>