Running stats handling on inference flow

GeorgeARM: <p>Hello all,</p>
<p>Have been trying to trace and export some vision models that contain operators like <code>BatchNorm2d</code> and end-up having all the running stats (e.g. <code>running_mean</code>/<code>running_var</code>) as buffer inputs instead of parameters.</p>
<p>Have had to manually remove them and handle with something like:</p>
<pre><code class="lang-auto">def clear_running_stats_for_inference(model):
    for _, module in model.named_modules():
        if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
            # Disable training semantics
            module.track_running_stats = False
            module.training = False

            # Remove batch tracking
            if hasattr(module, "num_batches_tracked"):
                delattr(module, "num_batches_tracked")

            # Set running_var/running_mean as parameters
            _buffer_to_parameter(module, "running_mean")
            _buffer_to_parameter(module, "running_var")

</code></pre>
<p>As I am new to this would like to understand if it would make sense for them to be handled like constants parameters during inference and what is the best way forward in order to avoid explicit handling.</p>

ptrblck: <p>Could you describe why you want to transform the buffers to parameters if you are not training them via gradient descent?</p>

GeorgeARM: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>From what I gather parameters are for trainable semantics; what I wanted to achieve is constant-ness when exporting for inference for these buffers. Moving them to parameters helped me achieve this.</p>
<p>To elaborate I am trying to export <code>mobilenet</code> and use <code>torch-mlir</code> to go down to the torch dialect; but I notice that all these stats end-up becoming inputs to the model as they are stored in buffers while they are really constant during inference.</p>
<p>Highly likely that this is something that can be handled at <code>torch-mlir</code> level really; trying to figure out.</p>

ptrblck: <aside class="quote no-group" data-username="GeorgeARM" data-post="3" data-topic="216063">
<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/georgearm/48/75389_2.png" class="avatar"> GeorgeARM:</div>
<blockquote>
<p>From what I gather parameters are for trainable semantics;</p>
</blockquote>
</aside>
<p>This is correct. <code>nn.Parameter</code>s are trainable and can be registered into <code>nn.Module</code>s.<br>
E.g. the <code>.weight</code> and <code>.bias</code> attributes of e.g. <code>nn.Linear</code> are trainable parameters:</p>
<pre data-code-wrap="python"><code class="lang-python">lin = nn.Linear(10, 10)
print(lin.weight)
# Parameter containing:
# tensor([[-0.2245,  0.2853, -0.2039,  0.1198,  0.1099,  0.2087,  0.1796, -0.1416,
#           0.2469, -0.2738],
#         ...

for name, param in lin.named_parameters():
    print(name)
# weight
# bias
</code></pre>
<aside class="quote no-group" data-username="GeorgeARM" data-post="3" data-topic="216063">
<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/georgearm/48/75389_2.png" class="avatar"> GeorgeARM:</div>
<blockquote>
<p>what I wanted to achieve is constant-ness when exporting for inference for these buffers. Moving them to parameters helped me achieve this.</p>
</blockquote>
</aside>
<p>In contrast to parameters, buffers are also registered to modules but are not trainable.<br>
The running stats are buffers since they are not trained via gradient descent and no gradient will be computed for them.</p>
<pre data-code-wrap="python"><code class="lang-python">bn = nn.BatchNorm2d(3)
for name, param in bn.named_parameters():
    print(name)
# weight
# bias
    
for name, buffer in bn.named_buffers():
    print(name)
# running_mean
# running_var
# num_batches_tracked
</code></pre>
<p>Transforming the <code>running_*</code> stats into parameters will make them trainable while you claim you want to achieve “constantness”, so I assume you want to keep them static.</p>

bdhirsh: <p>In addition to Piotr’s comments- is the problem you’re trying to solve the fact that the batchnorm buffers are showing up as graph inputs in the exported graph, but you’d like them to be constant s on the graph instead?</p>
<p>If you call .module() on the output of export, any module state (both params and buffers) will be <em>unlifted</em> back into state. You can do something like:</p>
<p>(returns a torch.fx.GraphModule, where any params/buffers that had previously been lifted into graph inputs will once again become module state)</p>
<p>torch.export.export(model).run_decompositions().module()</p>