Does torch.jit.script requires no_grad and eval modes?

mfoglio: <p>My goal is to export a model for fast inference.<br>
When I trace a model with <code>jit.trace</code> I usually do it this way:</p>
<pre><code class="lang-auto">model.eval()
with torch.no_grad():
    input = torch.rand(size=(1, 3, 500, 500))
    traced_cell = torch.jit.trace(model.to("cpu"), (input))
</code></pre>
<p>It is not clear to me if <code>model.eval()</code> and <code>with torch.no_grad()</code> are required (or still suggested) when exporting a model with <code>jit.script</code>.<br>
Thank you!</p>

ptrblck: <p><code>model.eval()</code> will change the behavior of some modules (e.g. dropout layers will be disabled and batchnorm layers will use their running stats to normalize the data). <code>torch.jit.trace</code> does not capture any data-dependent control flow, i.e. the code path used by the input will only be captured and other inputs won’t take a different path based on e.g. <code>if</code> statements etc. Given that, it sounds right to use <code>model.eval()</code> before tracing (otherwise the dropout layer would be used with the same mask in each forward pass). I don’t know if disabling the gradient calculation is needed during tracing or could also be added later during the inference. In any case, you might also want to check <a href="https://pytorch.org/docs/stable/generated/torch.autograd.inference_mode.html?highlight=inference_mode#torch.autograd.inference_mode"><code>torch.autograd.inference_mode()</code></a> for your model deployment.</p>

mfoglio: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a>! How about <code>jit.script</code>?</p>

ptrblck: <p><code>torch.jit.script</code> is able to capture data-dependent control flow (e.g. the dropout masks would be randomly sampled in each step in case you leave the model in <code>.train()</code> mode). However, the common use case during inference is to use the <code>.eval()</code> mode, so you might call it nevertheless even if scripting the model. The advantage would be that your model could use other data-dependent control flow, such as conditions based on the shape of the input etc.</p>

mfoglio: <p>Thank you. Should I call <code>.eval()</code> before <code>jit.script</code> or after it? What about <code>torch.autograd.inference_mode</code>? I don’t understcand if the results of these operations “are saved” by <code>jit.script</code>.<br>
Thanks!</p>

ptrblck: <p><code>jit.script</code> should not capture training/eval mode or the <code>no_grad()</code> context, so you should be able to script the model and call <code>.eval()</code> as well as <code>inference_mode/no_grad</code> during the deployment.<br>
However, if you are seeing any issue with this, please let us know.</p>

mfoglio: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> . My question is actually the following: is there any benefit on calling a <code>.eval()</code> and enable <code>no_grad</code> on torch script model? Is it something I should after I load a file in order to gain performance, or is it useless?</p>

ptrblck: <p>Yes, if you are deploying call <code>inference_mode/no_grad</code> for performance gains.<br>
<code>model.eval()</code> will change the behavior of some layers, such as disabling dropout and using the running stats of batchnorm layers so it’s not a performance (speed) improvement, but used during the evaluation/test phase of the model.</p>

mfoglio: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a> ! I have a follow up question: I have been recently fine tuned an <code>efficientnet_b0</code> model. If, after training the model, I do inference using the model without calling <code>model.eval()</code> first, the output of the model is completely wrong. Is this expected? Thank you!</p>

ptrblck: <p>If you are not calling <code>model.eval()</code> e.g. all batchnorm layers will use the input batch stats to normalize the activation as described before, which might be noisy depending on the batch size. Also, dropout would be enabled but I wouldn’t expect to see a huge difference to the training run (besides the noisy updates of course).</p>