Traced/Scripted models do not produce same output as eager models when weights are loaded

numpee: <p>JIT traced/scripted models are expected to produce the same output as eager models when given the same output.<br>
This seems to be true when we use randomly initialized weights, but not true when weights are loaded to models.<br>
Below is a code snippet to reproduce the issue:</p>
<h3><a name="without-loading-weights-1" class="anchor" href="#without-loading-weights-1"></a>Without loading weights:</h3>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torchvision.models import resnet18, ResNet18_Weights

# Create eager, scripted, and traced models
eager_model = resnet18().cuda(0)
eager_model.eval()
script_model = torch.jit.script(eager_model).cuda()
script_model.eval()
trace_input = torch.randn(4, 3, 224, 224).cuda()
traced_model = torch.jit.trace(eager_model, trace_input).cuda()
traced_model.eval()

# Random input and feature extraction
x = torch.randn(16, 3, 224, 224).cuda()

with torch.no_grad():
    eager_out = eager_model(x)
    script_out = script_model(x)
    traced_out = traced_model(x)
    
    if not torch.allclose(eager_out, script_out):
        print(f'Scripted: {(eager_out - script_out).abs().sum()}')
    if not torch.allclose(eager_out, traced_out):
        print(f'Traced: {(eager_out - traced_out).abs().sum()}')
</code></pre>
<p>There is no output for the code above.</p>
<h3><a name="using-pre-trained-weights-2" class="anchor" href="#using-pre-trained-weights-2"></a>Using pre-trained weights</h3>
<p>Now consider the code below:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torchvision.models import resnet18, ResNet18_Weights

# Create eager, scripted, traced models
eager_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1).cuda(0)
eager_model.eval()
script_model = torch.jit.script(eager_model).cuda()
script_model.eval()
trace_input = torch.randn(4, 3, 224, 224).cuda()
traced_model = torch.jit.trace(eager_model, trace_input).cuda()
traced_model.eval()

# Random input for feature extraction
x = torch.randn(16, 3, 224, 224).cuda()

with torch.no_grad():
    eager_out = eager_model(x)
    script_out = script_model(x)
    traced_out = traced_model(x)
    
    if not torch.allclose(eager_out, script_out):
        print(f'Scripted: {(eager_out - script_out).abs().sum()}')
    if not torch.allclose(eager_out, traced_out):
        print(f'Traced: {(eager_out - traced_out).abs().sum()}')
</code></pre>
<p>The output of the code is as follows:</p>
<blockquote>
<p>Scripted: 4.249152183532715<br>
Traced: 4.250102996826172</p>
</blockquote>
<p>Any ideas what’s going on? Is this expected behavior? A difference of ~4.25 is small considering that the output is a tensor of [16 x 1000], but not insignificant. Obviously for classification this shouldn’t be much of an issue, but I do suspect that in other cases (where precision is critical) it may pose issues.</p>
<p>BTW I am using <code>torch==1.13.1</code> and <code>torchvision==0.14.1</code></p>

ptrblck: <p>I cannot reproduce the issue deterministically in the first iteration, but after a few warmup iterations the <code>abs().max()</code> error increases. I assume this is caused by model optimizations, such as layer fusions etc., which are not creating bitwise-identical outputs. I also don’t see a difference between the random or pre-trained model in <code>torch==2.0.0</code>.</p>

numpee: <p>I also noticed something similar during my testing. Even when using <code>torch.ones(4, 3, 224, 224).cuda()</code> as the input, in many cases the first iteration would return identical outputs for all models, while any subsequent iterations would produce a certain amount of error. Again, I’m not sure if this is expected behavior — Does the compiled model still undergo optimizations (e.g., layer fusions as you mentioned) <em>after</em> compilation?</p>

ptrblck: <p>Yes, small errors caused by the limited numerical precision are expected, since the different algorithms will be used, i.e. especially if layers/operations are fused.</p>
<aside class="quote no-group" data-username="numpee" data-post="3" data-topic="176243">
<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/numpee/48/43681_2.png" class="avatar"> numpee:</div>
<blockquote>
<p>Does the compiled model still undergo optimizations (e.g., layer fusions as you mentioned) <em>after</em> compilation?</p>
</blockquote>
</aside>
<p>You are not using the newly added <code>torch.compile</code> approach, which would optimize the model directly (if I’m not mistaken), but you are scripting/tracing the model. The latter approach uses TorchScript and uses the first (3?) iterations to optimize the model.</p>

numpee: <p>I see. It just seems a bit odd that numerical precision issues will show up when weights are loaded, but not when models are randomly initialized.</p>
<p>I am actually looking into <code>torch.compile</code> in version 2.0. However, looking at the official blog post, it seems like the compiled model can’t be “exported” like the traced/scripted models (under the “Serialization” section) . The main reason I want to use traced/scripted models is so that I can train and export a model from one codebase and load it from another without having to copy the model code and call <code>load_state_dict()</code>. Please let me know if I’ve misunderstood this <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

siamak_z: <p>I have the same issue with PyTorch 1.12.1 and  CUDA 11.3. The traced version of my model with initialized random weights produces outputs closely matching the python model. However, when I load the weights from a checkpoint, the output of the traced model significantly (&gt;2.0) differs from that of the python model!<br>
<a class="mention" href="/u/ptrblck">@ptrblck</a> Any idea how this problem can be avoided?</p>