Is it good/neutral/bad practice to specify loss and optimizer inside the model class?

username: Arturas_Druteika
cooked: <p>I don’t know if it is even worth to ask this question in this forum, but I was wondering if it is a good, neutral or bad practice to specify loss as well as optimizer inside the model class, for example:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, in_features, n_classes, lr):
        self.fc_1 = nn.Linear(in_features, 256)
        self.fc_2 = nn.Linear(256, n_classes)

        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = nn.Adam(self.parameters(), lr=0.01)

    def forward(self, x)
        return self.fc_2(F.relu(self.fc_1(x)))
</code></pre>
<p>As of right now I’ve seen only one tutorial that uses this kind of approach. At first, for me, it looked kinda sketchy, but now I think that this is way more clean than creating criterion and optimizer outside of class.<br>
But I have a question, does this way have any implications or unnecessary behavior which I don’t understand or that I haven’t encountered that could harm my model’s performance or something like that?</p>

username: KFrank
cooked: <p>Hi Arturas!</p>
<aside class="quote no-group" data-username="Arturas_Druteika" data-post="1" data-topic="127749" data-full="true">
<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/arturas_druteika/48/53210_2.png" class="avatar"> Arturas_Druteika:</div>
<blockquote>
<p>if it is a good, neutral or bad practice to specify loss as well as optimizer inside the model class</p>
</blockquote>
</aside>
<p>Including your loss and/or optimizer in your model class won’t break<br>
anything or cause inefficiency.  So from a mechanical or technical<br>
perspective, it’s okay.</p>
<p>However, I would say that doing so is bad practice from a stylistic<br>
or design perspective.</p>
<p>The main idea is “separation of concerns.”  To my mind the purpose<br>
of a model is to take an input and produce a prediction.  At inference<br>
time (in contrast to training), you won’t be using a loss function nor<br>
optimizer, so they don’t really belong in the model class.  In a similar<br>
vein, when performing evaluation, you might choose to use a different<br>
“loss” function (perhaps a non-differentiable accuracy function) than<br>
the one you use for training.  Or you might use more than one accuracy<br>
metric.</p>
<p>It is also true that an important purpose of a model is to be able to be<br>
trained – that’s why models have <code>Parameter</code>s that have <code>grad</code>s and<br>
layers that have <code>backward()</code> functions.  It seems to me to be a sound<br>
stylistic / design choice to have these be part of the internals of the<br>
model class.</p>
<p>But the details of <em>how</em> to train the model (e.g., what optimization<br>
algorithm) can be separated from the infrastructure that makes a<br>
model able to be trained (e.g., tensors that have a <code>grad</code> property),<br>
and the guideline of “separation of concerns” suggests that when<br>
things can be conceptually and straightforwardly separated, they<br>
should be separated and placed in separate classes.</p>
<p>So “clean design” (whatever that means) suggests that the loss and<br>
optimizer aren’t “naturally” part of the model, so they shouldn’t be<br>
part of the model class.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Arturas_Druteika
cooked: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="127749">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>But the details of <em>how</em> to train the model (e.g., what optimization<br>
algorithm) can be separated from the infrastructure that makes a<br>
model able to be trained (e.g., tensors that have a <code>grad</code> property),<br>
and the guideline of “separation of concerns” suggests that when<br>
things can be conceptually and straightforwardly separated, they<br>
should be separated and placed in separate classes.</p>
</blockquote>
</aside>
<p>This is where I had doubts, because I thought that they are not part of the model, but at the same time, the example I gave made my feel that this is OK to do it like that. But now that I’ve heard another person’s opinion, I have to agree with you, that yes, loss and optimizer are not part of the model, and they shouldn’t be specified inside the model.<br>
But still, it would be interesting to see what more people have to say.</p>

username: ptrblck
cooked: <p>I agree with <a class="mention" href="/u/kfrank">@KFrank</a> also also think that the model, criterion, and optimizer are separate units, which can be independently used and changed.<br>
E.g. reusing your “model” to finetune it in another scenario would also recreate your optimizer, which would not be needed. Restoring a checkpoint also would need some additional check to see where the optimizer is located and used. In case you are using utilities, such as mixed-precision training via <code>torch.cuda.amp</code>, you would also need to manipulate the <code>forward</code> method of your model to add e.g. the <code>GradScaler</code>, and would need to add the “amp logic” to the model itself, which I would also like to avoid.</p>