Correct way to freeze a part of a model

yiftach: <p>Hi,<br>
I have a (outer) model that contains a (inner) backbone. Normally, they both train. Sometimes, I want to freeze the backbone. As far as I understand, this means:</p>
<ol>
<li>Once at the beginning - iterate over all parameters and set their <code>requires_grad</code> to <code>False</code></li>
<li>Make sure that the model is always set to <code>.eval()</code> and not <code>.train()</code>, to make sure it does not do dropout etc.<br>
The second part is what I’m concerned with. What’s the correct way to set the inner model to permanently be on <code>.eval()</code>, such that calling <code>.train()</code> on the outer model does not undo it? e.g.</li>
</ol>
<pre><code class="lang-auto">import torch
from torch import nn

class Inner(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc = nn.Linear()

class Outer(nn.Module):
    def __init__(self):
        super().__init__()

        self.inner = nn.Module()

outer = Outer()
outer.inner.eval()
print(outer.inner.training) # is False as desired

outer.train()
print(outer.inner.training) # is True, not desired
</code></pre>
<p>I could override the <code>.train()</code> function of the inner model and make it a no-op, but that seems hacky. What’s the best practice?</p>

Karthik_Ganesan: <p>If you try <code>outer.inner.eval()</code> after <code>outer.train()</code>, does that work?</p>

yiftach: <p>That works, but what I want is the cleanest option to mark a module as frozen and not have to think about it again (and having to remember calling <code>outer.inner.eval()</code> after each call to <code>outer.train()</code> is the opposite of that). I could also override <code>outer.train()</code> to always call <code>outer.inner.eval()</code> after it finishes running, but that also doesn’t seem very clean.</p>

Karthik_Ganesan: <p>The only option I can think of that would be ‘cleaner’ might be to define your own class member function(s) to do this. And inside that you can set the <code>.eval()</code> and <code>train()</code> for <code>inner</code> and <code>outer</code> as needed. Because you are doing something somewhat atypical, I don’t feel there would be any other way to do it, using just existing PyTorch functions.</p>

marksaroufim: <p><code>tensor.detach()</code> might be closer to what you’re looking for</p>

yiftach: <aside class="quote no-group" data-username="Karthik_Ganesan" data-post="4" data-topic="189298" 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/karthik_ganesan/48/34393_2.png" class="avatar"> Karthik_Ganesan:</div>
<blockquote>
<p>The only option I can think of that would be ‘cleaner’ might be to define your own class member function(s) to do this. And inside that you can set the <code>.eval()</code> and <code>train()</code> for <code>inner</code> and <code>outer</code> as needed. Because you are doing something somewhat atypical, I don’t feel there would be any other way to do it, using just existing PyTorch functions.</p>
</blockquote>
</aside>
<p>Thanks, but that removes some of the abstraction - the user can no longer use the outer model without knowing about its parts. Ideally, everything would look the same whether inner is a 1. trainable model, 2. a frozen mode or 3. a function - but in your suggestion, as far as I understand, 2. requires different usage.</p>
<aside class="quote no-group" data-username="marksaroufim" data-post="5" data-topic="189298" 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/marksaroufim/48/36747_2.png" class="avatar"> marksaroufim:</div>
<blockquote>
<p><code>tensor.detach()</code> might be closer to what you’re looking for</p>
</blockquote>
</aside>
<p>I’m familiar with <code>tensor.detach()</code> in other cases, could you please elaborate how using it here helps?</p>