Optimizer inside model, good practice?

username: thomas
cooked: <p>Hello,</p>
<p>I saw today on <a href="https://github.com/philtabor/Youtube-Code-Repository/blob/master/ReinforcementLearning/PolicyGradient/actor_critic/actor_critic_continuous.py" rel="noopener nofollow ugc">this</a> repository the practice of putting <code>optimizer</code> and <code>device</code> inside the model:</p>
<pre><code class="lang-auto">import torch as T
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

class GenericNetwork(nn.Module):
    def __init__(self, alpha, input_dims, fc1_dims, fc2_dims,
                 n_actions):
        super(GenericNetwork, self).__init__()
        self.input_dims = input_dims
        self.fc1_dims = fc1_dims
        self.fc2_dims = fc2_dims
        self.n_actions = n_actions
        self.fc1 = nn.Linear(*self.input_dims, self.fc1_dims)
        self.fc2 = nn.Linear(self.fc1_dims, self.fc2_dims)
        self.fc3 = nn.Linear(self.fc2_dims, self.n_actions)
        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cuda:1')
        self.to(self.device)

    def forward(self, observation):
        state = T.tensor(observation, dtype=T.float).to(self.device)
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x
</code></pre>
<p>I felt that it was petty ingenuous, especially when you are in a case with several networks (it avoids to have two times more self variables)</p>
<p>I was wondering if this was a good, or a bad practice and if there is some drawbacks. I never saw this before, so i am guessing that there is a good reason ?</p>

username: ptrblck
cooked: <p>I personally don’t use this workflow, as I don’t think that an optimizer is necessarily bound to a model.<br>
Of course your use case might use a single model and optimizer and thus this workflow will certainly work for your use case, but it’s “cleaner” to me to use them as separate objects.</p>
<p>One potential drawback would also be the creation of checkpoints. While <code>model.state_dict()</code> would properly return all parameters and buffers, you would now need to access the <code>optimizer</code> object through the <code>model</code> to save and load its <code>state_dict</code>. Again, this comes down to your coding style, so I’m sure there might also be valid reasons to use this approach.</p>
<p>Also, I wouldn’t add the device information inside the model directly, as this would break e.g. <code>nn.DataParallel</code>.</p>

username: thomas
cooked: <p>Thanks for your answer ! It does seems indeed inconvenient for saving/loading<br>
I am not sure I did understand the <code>nn.DataParallel</code> part, is it possible to give a very small example or explanation ?</p>

username: ptrblck
cooked: <p>The issue for <code>nn.DataParallel</code> is that the model uses a hardcoded <code>device</code> as <code>cuda:0</code>, if CUDA is available (and <code>cuda:1</code>, which will fail so I assume <code>'cpu'</code> should be used in the <code>else</code> statement).<br>
However, <code>nn.DataParallel</code> will copy the model to each specified device in each forward pass, will split the input in <code>dim0</code>, and send each chunk to the corresponding device.<br>
For this to work you cannot use hardcoded device arguments inside the model or <code>forward</code>, since <code>nn.DataParallel</code> will take care of it. If you need to create new tensors inside the <code>forward</code>, you would have to use the currently used device for this model replica (e.g. via <code>new_tensor = torch.randn(1, device=x.device)</code>).</p>
<p>While this code works fine on the GPU:</p>
<pre><code class="lang-python">model = GenericNetwork(1., (1,), 1, 1, 1)
x = T.randn(8, 1)
out = model(x)
print(out)
</code></pre>
<p>You’ll get a device mismatch error, if you use <code>model = nn.DataParallel(model)</code>:</p>
<pre><code class="lang-python">RuntimeError: Expected tensor for 'out' to have the same device as tensor for argument #2 'mat1'; but device 1 does not equal 0 (while checking arguments for addmm)
</code></pre>
<p>for the aforementioned reason.</p>