Problem converting model to onnx

RealNewNoob: <p>I am working on this code to convert my Unet++ model to Onnx, and the inferenced results from pytorch and onnx have huge differences.</p>
<p>I don’t need any dynamic axis, they are just there since they are suggested by GPT somehow.</p>
<p>Enviroment, windows, pytorch torch  2.7.1+cu128, onnxruntime-gpu 1.22.0, onnx 1.19.1.</p>
<p>Model Implementation</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import nn

__all__ = ["UNet", "NestedUNet"]


class VGGBlock(nn.Module):
    def __init__(self, in_channels, middle_channels, out_channels):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(middle_channels)
        self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        return out


class UNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, **kwargs):
        super().__init__()

        nb_filter = [32, 64, 128, 256, 512]

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)

        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])

        self.conv3_1 = VGGBlock(nb_filter[3] + nb_filter[4], nb_filter[3], nb_filter[3])
        self.conv2_2 = VGGBlock(nb_filter[2] + nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv1_3 = VGGBlock(nb_filter[1] + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv0_4 = VGGBlock(nb_filter[0] + nb_filter[1], nb_filter[0], nb_filter[0])

        self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)

    def forward(self, input):
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))

        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))

        output = self.final(x0_4)
        return output


class NestedUNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, deep_supervision=False, **kwargs):
        super().__init__()

        nb_filter = [32, 64, 128, 256, 512]

        self.deep_supervision = deep_supervision

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)

        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])

        self.conv0_1 = VGGBlock(nb_filter[0] + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_1 = VGGBlock(nb_filter[1] + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_1 = VGGBlock(nb_filter[2] + nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv3_1 = VGGBlock(nb_filter[3] + nb_filter[4], nb_filter[3], nb_filter[3])

        self.conv0_2 = VGGBlock(
            nb_filter[0] * 2 + nb_filter[1], nb_filter[0], nb_filter[0]
        )
        self.conv1_2 = VGGBlock(
            nb_filter[1] * 2 + nb_filter[2], nb_filter[1], nb_filter[1]
        )
        self.conv2_2 = VGGBlock(
            nb_filter[2] * 2 + nb_filter[3], nb_filter[2], nb_filter[2]
        )

        self.conv0_3 = VGGBlock(
            nb_filter[0] * 3 + nb_filter[1], nb_filter[0], nb_filter[0]
        )
        self.conv1_3 = VGGBlock(
            nb_filter[1] * 3 + nb_filter[2], nb_filter[1], nb_filter[1]
        )

        self.conv0_4 = VGGBlock(
            nb_filter[0] * 4 + nb_filter[1], nb_filter[0], nb_filter[0]
        )

        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)

    def forward(self, input):
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))

        x2_0 = self.conv2_0(self.pool(x1_0))
        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))

        x3_0 = self.conv3_0(self.pool(x2_0))
        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))

        if self.deep_supervision:
            output1 = self.final1(x0_1)
            output2 = self.final2(x0_2)
            output3 = self.final3(x0_3)
            output4 = self.final4(x0_4)
            return [output1, output2, output3, output4]

        else:
            output = self.final(x0_4)
            return output
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">import torch
import numpy as np
from archs import NestedUNet, UNet
from pathlib import Path
import os
import onnxruntime as ort

experiment_dir = Path("models/wafer_defect_dataset_0918_NestedUNet_woDS_0")
load_model_path = str(experiment_dir/"model.pth")
export_model_path = str(experiment_dir/f"model.onnx")

if Path(export_model_path).exists():
    print(f"模型已存在: {export_model_path}")
    os.remove(export_model_path)
    print(f"删除模型: {export_model_path}")

# 加载模型并迁移到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device != torch.device("cuda"):
    raise ValueError("GPU不可用")

if "NestedUNet" in load_model_path:
    model = NestedUNet(num_classes=1, input_channels=3, deep_supervision=False)
else:
    model = UNet(num_classes=1, input_channels=3, deep_supervision=False)
model.load_state_dict(torch.load(load_model_path, map_location=device))  # 加载权重时指定设备
model.to(device)
model.eval()

# 生成固定随机种子的输入（确保可复现）
torch.manual_seed(42)
dummy_input = torch.randn(1, 3, 512, 512, device=device)  # 直接在目标设备生成输入

with torch.no_grad():
    pytorch_output = model(dummy_input).cpu().numpy()

# 导出ONNX（禁用动态轴，避免维度相关优化导致的差异）
torch.onnx.export(
    model,
    (dummy_input,),
    export_model_path,
    input_names=["input"],
    output_names=["output"],
    opset_version=12,                 # 建议 13
    do_constant_folding=False,        
    # dynamic_axes=None,
    dynamic_axes={
    "input": {2: "height", 3: "width"},
    "output": {2: "height", 3: "width"},
    },
    verbose=False,
)

# 验证PyTorch/ONNX输出（CPU/CUDA，禁优化）
def run_and_diff(provider: str):
    try:
        so = ort.SessionOptions()  # type: ignore[attr-defined]
        try:
            so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL  # type: ignore[attr-defined]
        except Exception:
            pass
        session = ort.InferenceSession(str(export_model_path), sess_options=so, providers=[provider])  # type: ignore[attr-defined]
    except Exception:
        session = ort.InferenceSession(str(export_model_path), providers=[provider])  # type: ignore[attr-defined]
    onnx_input = {session.get_inputs()[0].name: dummy_input.cpu().numpy()}
    onnx_output = session.run(None, onnx_input)[0]
    diff = np.abs(pytorch_output - onnx_output)
    print(f"[{provider}] 最大差异: {diff.max():.6f}")
    print(f"[{provider}] 平均差异: {diff.mean():.6f}")

run_and_diff("CPUExecutionProvider")
if torch.cuda.is_available():
    try:
        run_and_diff("CUDAExecutionProvider")
    except Exception as e:
        print(f"CUDAExecutionProvider 运行失败: {e}")
</code></pre>
<p>Terminal output<br>
模型已存在: models\wafer_defect_dataset_0918_NestedUNet_woDS_0\model.onnx<br>
删除模型: models\wafer_defect_dataset_0918_NestedUNet_woDS_0\model.onnx<br>
[CPUExecutionProvider] 最大差异: 2.782166<br>
[CPUExecutionProvider] 平均差异: 0.457833<br>
2025-10-16 19:34:30.1944079 [W:onnxruntime:, transformer_memcpy.cc:83 onnxruntime::MemcpyTransformer::ApplyImpl] 10 Memcpy nodes are added to the graph main_graph for CUDAExecutionProvider. It might have negative impact on performance (including unable to run CUDA graph). Set session_options.log_severity_level=1 to see the detail logs before this message.<br>
[CUDAExecutionProvider] 最大差异: 2.393982<br>
[CUDAExecutionProvider] 平均差异: 0.202145</p>
<p>Anyone please help me！ <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">  <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=14" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"></p>

RealNewNoob: <p>The problem was resolved. It was caused by the model weights, the variance of the normalization layer is very small due to errors during training (normalizing).</p>