[Performance] Outputs are completely different between .onnx model and .pth model

[timcov]

Describe the issue

So i trained a resnet18 on patched images. I run the resnet (.pth) file on my unique test set and get good results. I try to export it to onnx and i get way worse results.

I've checked the input tiles and they're identical, same max and min, same pre processing. I even exported the raw tiles and they were identical. I even looked at the individual pixel values and these were also the same.

Now when i run the image trough the onnx model the export logits are different.

(Same result on both windows and ubuntu)

Cam_1_Frame_325_HasKnots_False_ts_2025-11-03-11-29-16-579.tiff
Cam_1_Frame_355_HasKnots_True_ts_2025-11-03-14-14-04-199.tiff

Normally all the files should be here, if something doesn't work. Be sure to let me know because i changed it a bit so i wouldnt have to upload everything

config.py
export_to_onnx.py
test_onnx.py
test_pth.py

To reproduce

• uv add onnxruntime
• uv add onnx
• uv pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu128
• uv add onnxscript
• uv add scikit-learn

Then run the export to onnx as follows:
uv run python export_to_onnx.py --model "resnet18" --weights "./resnet18_best.pth" --output "test.onnx"

Then
The root folder should contain 2 subfolders 1.Knot and 0.NoKnot. You can just copy the same image twice

uv run python test_onnx.py --onnx "./test.onnx" --root ./images

I also added 2 images, the first one is without knots and the logits of this image should be
[-3.072449 -2.8448505 -7.6250186 -6.501796 ] -> is like this for .pth file but not for .onnx file

The second one has knots and he should return 3 positive classe (1) and 1 negative (0)

Urgency

It is quite urgent as i need it within this week.

Platform

Windows

OS Version

Windows 11

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

onnxruntime 1.23.2

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

Cuda 12.8 and 12.6

Model File

If needed i can try to upload it.

Is this a quantized model?

No

[justinchuby]

Please test with torch.onnx.export(..., dynamo=True, report=True) using the latest torch-nightly. Attach the generated report if there is an error. Thanks!

You may also need the latest onnxscript build from the main branch because of microsoft/onnxscript#2680

[timcov]

Thanks for the quick response @justinchuby. I ran it and there are no errors

No errors
Besides this i also added the latest torch nightly and onnxscript

torch 2.10.0.dev20251105+cu126
torchvision 0.25.0.dev20251105+cu126
onnxscript 0.5.6

But the issue seems to continue.

There is however 1 small difference between the tensors sent to the onnx and the pth. The .pth version is rounded off a bit earlier. 
I dont know if this is a big issue or not.

PTH:
tensor([[[[-1.9124, -1.8953, -1.8953,  ..., -1.6213, -1.6213, -1.6555],
[[[[-1.9124069 -1.8952821 -1.8952821 ... -1.6212862 -1.6212862 -1.6555357]
     [-1.9124, -1.8953, -1.8953,  ..., -1.6555, -1.6555, -1.6384],
   [-1.9124069 -1.8952821 -1.8952821 ... -1.6555357 -1.6555357 -1.6384109]
     [-1.9124, -1.8953, -1.8953,  ..., -1.6555, -1.6555, -1.6384]
...

ONNX:
[[[[-1.9124069 -1.8952821 -1.8952821 ... -1.6212862 -1.6212862 -1.6555357]
          [-1.9124, -1.8953, -1.8953,  ..., -1.6555, -1.6555, -1.6384],
   [-1.9124069 -1.8952821 -1.8952821 ... -1.6555357 -1.6555357 -1.6384109]
          [-1.9124, -1.8953, -1.8953,  ..., -1.6555, -1.6555, -1.6384],
   [-1.9124069 -1.8952821 -1.8952821 ... -1.6555357 -1.6555357 -1.6384109] 
...

[timcov]

Now i also tried to use the exact same preprocessing, so i changed my onnx model to just do the following:

  def forward(self, tiles):
      logits = self.model(tiles)
      return logits, logits

Then i added it right before my .pth model in my test

with torch.no_grad():
    for patches, label, path in test_loader:
        sess = create_session(onnx_path)
        in_name = sess.get_inputs()[0].name
        outputs = sess.run(None, {in_name: patches.cpu().numpy()})
        print(f"ONNX outputs: {outputs}")
        logits = model(patches) 

And for the same input there are 2 different outputs

ONNX outputs: [array([[-2.9795594],
[-3.2486563],
[-2.2178779],
[-2.025703 ]], dtype=float32)

Logits per patch: [-3.072449 -2.8448505 -7.6250186 -6.501796 ] (PTH)

Could the difference be GPU for pth and CPU for onnx?

[timcov]

Update:

I exported another older checkpoint and it also gives different results than the earlier onnx version (Same code). This makes me think that it might be a version mismatch somewhere. I'll look futher into it if its something else.

(non patched version)

Edit: It is almost as the models differ, but the exact same .pth file was used. Or that the weights are incorrectly parsed/binded to the onnx model

[timcov]

This seemed to fix it

with torch.inference + dynamo False + training=torch.onnx.TrainingMode.EVAL + constant folding and keep_initializers_as_inputs=False

"""Export a quadrant-patching model to ONNX format."""

import argparse
import os
import sys
import torch
from torch import nn

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from models import load_model  # your existing loader

# export_model_to_onnx_quadrants.py  (only the wrapper shown)

class ONNXWrapper(nn.Module):
    """
    NHWC uint8 full image(s) -> TL/TR/BL/BR tiles -> resize to (target_h, target_w)
    -> normalize (ImageNet) -> model -> per-tile probs [4B,2] & preds [4B,1].
    """
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.register_buffer("mean_imagenet", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std_imagenet",  torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, x):
        """
        Splits input image into quadrants, preprocesses, and runs through the model.
        Args:
            image (torch.Tensor): Input image tensor in NHWC uint8 format.
        Returns:
            tuple: Preprocessed tiles and predictions.
        """

        tiles = []
        tile_h, tile_w = 1024, 1224  # hardcoded for quadrant tiling
        for i in range(2):
            for j in range(2):
                start_h = i * tile_h
                end_h = (i + 1) * tile_h
                start_w = j * tile_w
                end_w = (j + 1) * tile_w
                tile = x[:, start_h:end_h, start_w:end_w, :] # [B, tile_h, tile_w, 3]
                tiles.append(tile)

        tiles = torch.cat(tiles, dim=0)  # (4, tile_h, tile_w, 3)
        tiles = tiles.permute(0, 3, 1, 2).float()  # (4, 3, tile_h, tile_w)
        tiles = tiles / 255.0
        tiles = (tiles - self.mean_imagenet) / self.std_imagenet

        logits = self.model(tiles)  # [4, 1] or [4, 2]
        if logits.ndim == 1:
            logits = logits.unsqueeze(1)
        if logits.ndim == 2 and logits.shape[1] == 1:
            conf = torch.sigmoid(logits)
            preds = (conf > 0.5).to(logits.dtype)
        else:
            conf = torch.sigmoid(logits)
            preds = torch.argmax(conf, dim=1)
        return conf, preds 

def export_to_onnx(model_name, weights_path, output_path,
                   dummy_full_w=2448, dummy_full_h=2048):
    """
    Exports the model to ONNX format using quadrant tiling.
    """
    device = torch.device("cpu")
    print(f"using device: {device}")

    model = load_model(model_name, 1224, 1024)
    model.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
    model.eval().cpu()

    wrapper = ONNXWrapper(model)

    dummy = torch.zeros(1, 2048, 2448, 3, dtype=torch.uint8)

    # torch.onnx.export(
    #     wrapper,
    #     dummy,
    #     output_path,
    #     input_names=["image"],
    #     output_names=["probabilities", "predictions"],
    #     dynamo=True,
    #     report=True,
    #     export_params=True,
    #     opset_version=18,
    #     do_constant_folding=True
    # )
    model.eval()
    with torch.inference_mode():
        torch.onnx.export(
            wrapper,                     
            dummy,                        
            output_path,
            opset_version=18,
            dynamo=False,                
            training=torch.onnx.TrainingMode.EVAL,
            do_constant_folding=True,
            keep_initializers_as_inputs=False,
            input_names=["input"],
            output_names=["probabilities", "predictions"],
        )

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Export quadrant-patching model to ONNX")
    ap.add_argument("--model", required=True, help="Model name (e.g. resnet18)")
    ap.add_argument("--weights", required=True, help="Path to .pth")
    ap.add_argument("--output", required=True, help="Path to .onnx")
    ap.add_argument("--width",  type=int, default=1224, help="Trained PATCH width")
    ap.add_argument("--height", type=int, default=1024, help="Trained PATCH height")
    # You can leave dummy size at 2448x2048; dynamic axes make runtime flexible
    ap.add_argument("--dummy_full_width",  type=int, default=2448)
    ap.add_argument("--dummy_full_height", type=int, default=2048)
    args = ap.parse_args()

    export_to_onnx(
        args.model, args.weights, args.output
    )