设置track_running_stats=False时，ONNX导出BatchNorm仍处于训练模式的问题

核心问题原因 #

当track_running_stats=False时，PyTorch的BatchNorm会在推理阶段计算当前输入的均值方差，但你导出ONNX时使用了training=torch.onnx.TrainingMode.PRESERVE，这会保留模型的训练模式特征，导致ONNX将部分BatchNorm层标记为训练模式。另外，直接设置running_mean=None的方式没有完全同步模型的内部状态。

具体解决步骤 #

1. 修改ONNX导出的训练模式参数 #

将导出时的training参数从PRESERVE改为EVAL，强制模型以推理模式导出：

dummy_input = torch.randn(1, 3, 200, 200, requires_grad=True)  
torch.onnx.export(
    model, 
    dummy_input, 
    model_path,  
    export_params=True, 
    opset_version=16, 
    training=torch.onnx.TrainingMode.EVAL  # 修改为推理模式
)

2. 优化BatchNorm状态设置逻辑 #

使用named_modules()遍历所有子模块（避免children()漏过嵌套的BatchNorm），同时调用reset_running_stats()确保状态重置生效：

model.eval()

for name, module in model.named_modules():
    if isinstance(module, nn.BatchNorm2d):
        module.track_running_stats = False
        module.reset_running_stats()  # 重置内部缓存，确保forward逻辑正确

3. （可选）手动修正已导出的ONNX模型 #

如果上述方法无效，可以直接修改已导出的ONNX文件，将BatchNorm的training_mode从1（训练模式）改为0（推理模式）：

import onnx

onnx_model = onnx.load(model_path)
for node in onnx_model.graph.node:
    if node.op_type == "BatchNormalization":
        for attr in node.attribute:
            if attr.name == "training_mode":
                attr.i = 0
# 保存修正后的模型
onnx.save(onnx_model, "fixed_" + model_path)

验证方法 #

导出后重新检查ONNX节点，确认所有BatchNorm的training_mode=0，且多余的ReduceMean节点被优化掉，再用OpenVINO加载验证是否报错。

内容的提问来源于stack exchange，提问作者lpe