Get different results after JIT compilation optimization

tinywisdom: <p>I have a simple model file as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.fx as fx

class CustomTracer(fx.Tracer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def is_leaf_module(self, module, name):
        return isinstance(module, CustomModule)

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

    def forward(self, x):
        conv = nn.Conv2d(3, 64, 3, padding=1)
        return conv(x)

def utility_function(x):
    return torch.relu(x)

class UtilityClass:
    def __init__(self):
        pass

    def method(self, x):
        bn = nn.BatchNorm2d(x.size(1))
        return bn(x)

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.custom_module = CustomModule()
        self.utility_class = UtilityClass()

    def forward(self, x):
        x = self.custom_module(x)
        x = utility_function(x)
        x = self.utility_class.method(x)
        return x

def my_model_function():
    return MyModel()
</code></pre>
<p>I get my model like this:<code>model_function = getattr(model_module, "my_model_function")</code>.</p>
<p>I use two methods to derive the results of the model:</p>
<ul>
<li><code>output_original = model(input)</code></li>
<li><code>python jit_model = torch.jit.trace(model, input)    output_jit = jit_model(input)</code></li>
</ul>
<p>and get this output:</p>
<pre><code class="lang-auto">Mismatched elements: 3211246 / 3211264 (100.0%)
Greatest absolute difference: 60.796468168497086 at index (0, 22, 5, 223) (up to 1e-05 allowed)
Greatest relative difference: 19350106.104020353 at index (0, 16, 223, 168) (up to 1e-05 allowed)
</code></pre>
<p>Why this happened?</p>

ptrblck: <p>Your model does not even produce the same outputs in Eager mode:</p>
<pre data-code-wrap="python"><code class="lang-python">model = my_model_function()
x = torch.randn(1, 3, 24, 24)
out1 = model(x)
out2 = model(x)

torch.testing.assert_allclose(out1, out2)
# Mismatched elements: 36863 / 36864 (100.0%)
# Greatest absolute difference: 7.237934112548828 at index (0, 11, 11, 19) (up to 1e-05 allowed)
# Greatest relative difference: 49341.4609375 at index (0, 50, 23, 10) (up to 0.0001 allowed)
</code></pre>
<p>so I would probably focus on this first before digging into the deprecated <code>jit.trace</code> util.</p>
<p>The differences are most likely caused by the module creation in the <code>forward</code> pass, so move them into the <code>__init__</code> of the parent class.</p>

tinywisdom: <p>Ah, I see!!  Thanks for your help.</p>
<p>And, hi, sorry to bother you! I’m just curious — may I ask why you said <code>torch.jit.trace</code> is “deprecated” or not recommended? I can still find it in the official PyTorch 2.6 documentation, and it seems to be supported. Would you mind sharing any official notes, issues, or discussions that mention its deprecation or reduced support? I’d really appreciate it. Thank you!</p>

ptrblck: <p>You can find a few comments on GitHub and e.g <a href="https://github.com/pytorch/pytorch/issues/103841#issuecomment-1605017153">this one</a> from June 2023 explains the feature development is already “frozen”.</p>