Results are different when saving in TorchScript Format through train/eval mode

li.hao: <p>Hi community. Here is the <a href="https://github.com/pytorch/pytorch/issues/71929" rel="noopener nofollow ugc">issue report</a> I created. Can someone provide some insight or feedback in it? Thanks!</p>
<h3>
<a name="describe-the-bug-1" class="anchor" href="#describe-the-bug-1"></a><img src="https://discuss.pytorch.org/images/emoji/apple/bug.png?v=12" title=":bug:" class="emoji" alt=":bug:" loading="lazy" width="20" height="20"> Describe the bug</h3>
<p>Seems like the dropout layer is not disabled when saving a <code>Module</code> in <code>train</code> mode through <code>jit</code>.</p>
<h3>
<a name="reproduce-2" class="anchor" href="#reproduce-2"></a>Reproduce</h3>
<p>Creating and saving the model via:</p>
<ul>
<li>states</li>
<li>JIT script (with model in train mode)</li>
<li>JIT script (with model in eval mode)</li>
</ul>
<pre><code class="lang-python">import torch
import torch.nn as nn
import numpy as np

torch.manual_seed(1)
np.random.seed(1)
x = torch.tensor(np.random.random(100).astype(np.float32).reshape(1, 1, 10, 10))
net = nn.Sequential(nn.Flatten(), nn.Linear(100, 4), nn.Dropout(0.5))
# set to eval mode
net.eval()
print(net(x).tolist())

# saving states
torch.save(net.state_dict(), "/tmp/saved_states.pth")

# saving jit script in train mode
net.train()
traced_script_module = torch.jit.trace(net, x)
traced_script_module.save("/tmp/traced_train.pt")

# saving jit script in eval mode
net.eval()
traced_script_module = torch.jit.trace(net, x)
traced_script_module.save("/tmp/traced_eval.pt")
</code></pre>
<p>Loading back model in three approaches and evaluating them.</p>
<pre><code class="lang-python">reloaded_net_states = nn.Sequential(nn.Flatten(), nn.Linear(100, 4), nn.Dropout(0.5))
states = torch.load("/tmp/saved_states.pth")
reloaded_net_states.load_state_dict(states)

loaded_nets = {
    "states": reloaded_net_states,
    "traced_train": torch.jit.load("/tmp/traced_train.pt"),
    "traced_eval": torch.jit.load("/tmp/traced_eval.pt"),
}

for name, loaded_net in loaded_nets.items():
    print("testing", name)
    for i in range(3):
        # set to eval mode
        loaded_net.eval()
        print(loaded_net(x).tolist())
</code></pre>
<h3>
<a name="output-3" class="anchor" href="#output-3"></a>Output</h3>
<pre><code class="lang-auto">[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
testing states
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
testing traced_train
[[-0.20582711696624756, 0.0, -0.0, -0.0]]
[[-0.20582711696624756, 0.7365781664848328, -0.0, -0.3352513909339905]]
[[-0.0, 0.7365781664848328, -0.0, -0.0]]
testing traced_eval
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
[[-0.10291355848312378, 0.3682890832424164, -0.04384330287575722, -0.16762569546699524]]
</code></pre>
<h3>
<a name="expected-output-4" class="anchor" href="#expected-output-4"></a>Expected Output</h3>
<p>The results of saved jit scripts should be the same.</p>
<h3>
<a name="versions-5" class="anchor" href="#versions-5"></a>Versions</h3>
<p>Collecting environment information…<br>
PyTorch version: 1.10.0+cu102<br>
Is debug build: False<br>
CUDA used to build PyTorch: 10.2<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Ubuntu 20.04.3 LTS (x86_64)<br>
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0<br>
Clang version: Could not collect<br>
CMake version: version 3.16.3<br>
Libc version: glibc-2.31</p>
<p>Python version: 3.9.7 | packaged by conda-forge | (default, Sep 29 2021, 19:20:46)  [GCC 9.4.0] (64-bit runtime)<br>
Python platform: Linux-5.11.0-38-generic-x86_64-with-glibc2.31<br>
Is CUDA available: True<br>
CUDA runtime version: Could not collect<br>
GPU models and configuration:<br>
GPU 0: NVIDIA TITAN RTX<br>
GPU 1: NVIDIA TITAN RTX</p>
<p>Nvidia driver version: 495.29.05<br>
cuDNN version: Probably one of the following:<br>
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.0<br>
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.0<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A</p>
<p>Versions of relevant libraries:<br>
[pip3] numpy==1.21.4<br>
[pip3] torch==1.10.0<br>
[pip3] torchaudio==0.10.0+cu113<br>
[pip3] torchvision==0.11.1+cu113<br>
[conda] numpy                     1.21.4           py39hdbf815f_0    conda-forge<br>
[conda] torch                     1.10.0+cu113             pypi_0    pypi<br>
[conda] torchaudio                0.10.0+cu113             pypi_0    pypi<br>
[conda] torchvision               0.11.1+cu113             pypi_0    pypi</p>
<p>PS: I also reproduced it on version <code>1.9.0</code></p>

ptrblck: <p>Answered <a href="https://github.com/pytorch/pytorch/issues/71929#issuecomment-1064842494">here</a>.<br>
TL;DR: tracing a module records the actual execution (with or without dropout) and you thus cannot change the behavior afterwards.</p>

li.hao: <p>Thank you so much for your help. I will use scripting instead of tracing.</p>