Calling .eval() on a torch.jit.ScriptModule leaves paramers.requires_grad True

alekseynp: <p>I was having trouble with memory usage of my traced model in C++, and I discovered that <code>.eval()</code> doesn’t change the <code>requires_grad</code> for the parameters in my ScriptModule. Is this intended behaviour? I can say that as a user this was not expected behaviour. As a user I would like it to work as I expect, to warn, or to raise. Given that I can manually set the requires_grad behaviour, it seems like my expected behaviour is possible?</p>
<p>I think the underlying cause is that <code>my_script_module.layer</code> is a <code>RecursiveScriptModule</code> and has no <code>.children()</code>.</p>
<p>PyTorch 1.5</p>
<pre><code class="lang-auto">import torch


class MyScriptModule(torch.jit.ScriptModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(1, 1, bias=False)


my_script_module = MyScriptModule()

# [True]
print([p.requires_grad for p in my_script_module.parameters()])

my_script_module.eval()

# [True] :(
print([p.requires_grad for p in my_script_module.parameters()])

for p in my_script_module.parameters():
    p.requires_grad = False

# [False] :)
print([p.requires_grad for p in my_script_module.parameters()])
</code></pre>
<hr>
<p>I know this isn’t a totally normal thing to be doing. For what it’s worth, I am subclassing <code>ScrtipModule</code> in this way so that I can do the following. Maybe I should do something differently?</p>
<pre><code class="lang-auto">class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.inner = MyScriptModule()

    def forward(self, x):
        # stuff
        x = self.inner(x)
        # more stuff
        return x


class MyScriptModule(torch.jit.ScriptModule):
    """ Psuedo-code """
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(1, 1, bias=False)

    @torch.jit.script_method
    def forward(self, x):
        out = torch.zeros_like(x)
        for i in range(x.size()[0]):
            out = self.layer(x)
        return out


my_module = MyModule()
my_module.eval()

torch.jit.trace(my_module, sample)
</code></pre>

ptrblck: <aside class="quote no-group" data-username="alekseynp" data-post="1" data-topic="86791">
<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/alekseynp/48/26000_2.png" class="avatar"> alekseynp:</div>
<blockquote>
<p>I discovered that <code>.eval()</code> doesn’t change the <code>requires_grad</code> for the parameters in my ScriptModule. Is this intended behaviour?</p>
</blockquote>
</aside>
<p>Yes, this is expected and also won’t change the <code>requires_grad</code> attributes of an “eager” model.<br>
<code>model.eval()</code> and <code>model.train()</code> change the internal <code>self.training</code> flag of all modules recursively starting from the parent module. By doing so, the behavior of some layers will be changed.<br>
E.g. dropout will be disabled and batchnorm layers will use their running statistics to normalize the incoming data instead of calculating the batch statistics.</p>
<p>If you want to freeze the parameters, you would have to set their <code>.requires_grad</code> attribute to <code>False</code>.</p>
<p>You could of course use both in combination, e.g. freeze all parameters, but leave the dropout layers enabled, or let all parameters train, but use the running stats of batchnorm layers.</p>