ONNX Runtime giving a different result to Pytorch on a UNET output

[dmagee]

I'm having trouble getting a UNET running in ONNX runtime (in python) that works fine in PyTorch. The code I have is below (basically it loads a PT model, does inference on a single image, converts to ONNX, and then does inference on the same image with that.The outputs are different as shown by differning min/max/mean. I'm not sure if it's the oNNX conversion not working, or the way I'm calling the runtime.

import torch
from PIL import Image
from torchvision import transforms as T
import torchvision

def get_transform(image):
    transform = T.ToTensor()
    image = transform(image)
    return image


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Load image data
pic = Image.open('snapshot10X-crop.jpg')
pic=(get_transform(pic).unsqueeze(0)).to(device)


# Load PyTorch model
model_gc = torch.load("mymodel.pt")
model_gc.to(device)


print("In Tensor size:"+str(pic.size()))

# do inference in PyTorch (This is correct)
out = model_gc(pic)

print("Out Tensor size:"+str(out.size()))

# Calculate statistics of PyTorch output
max_value = torch.max(out[0,0,:,:])
min_value = torch.min(out[0,0,:,:])
mean_value = torch.mean(out[0,0,:,:])
print("Out0 Max,Min,Mean:")
print(max_value)
print(min_value)
print(mean_value)

max_value = torch.max(out[0,1,:,:])
min_value = torch.min(out[0,1,:,:])
mean_value = torch.mean(out[0,1,:,:])
print("Out1 Max,Min,Mean:")
print(max_value)
print(min_value)
print(mean_value)

#Save results as pngs (these look correct)
torchvision.io.write_png((out[0,0:1,:,:]*255).cpu().type(torch.uint8),"result0.png")
torchvision.io.write_png((out[0,1:2,:,:]*255).cpu().type(torch.uint8),"result1.png")

# Convert to onnx
from torch.onnx import export
export(model_gc, pic, "mymodel.onnx", opset_version=11, input_names=['image'] ,output_names=['output'])

# Now load, and do inference with ONNX model

import onnxruntime

session = onnxruntime.InferenceSession("mymodel.onnx")

# Check if the model has been loaded successfully
if session is None:
    raise ValueError("Failed to load the model")

result = session.run(["output"], {"image": pic.cpu().numpy()})

import numpy as np
out = result[0]

max_value = np.max(out[0,0,:,:])
min_value = np.min(out[0,0,:,:])
mean_value = np.mean(out[0,0,:,:])
print("ONNX Out0 Max,Min,Mean:")
print(max_value)
print(min_value)
print(mean_value)

max_value = np.max(out[0,1,:,:])
min_value = np.min(out[0,1,:,:])
mean_value = np.mean(out[0,1,:,:])
print("ONNX Out1 Max,Min,Mean:")
print(max_value)
print(min_value)
print(mean_value)

The model is at; https://drive.google.com/file/d/16nJfbSAzaIou7cnCamAeG8PQuCcyYpcT/view?usp=sharing
The image is at: https://drive.google.com/file/d/1aw75jZFNH98NJbTs2v9huVECGLEyMdMc/view?usp=sharing

Outputs are:

PyTorch:

Out0 Max,Min,Mean:
tensor(1., device='cuda:0', grad_fn=<MaxBackward1>)
tensor(0.0086, device='cuda:0', grad_fn=<MinBackward1>)
tensor(0.8373, device='cuda:0', grad_fn=<MeanBackward0>)
Out1 Max,Min,Mean:
tensor(0.9914, device='cuda:0', grad_fn=<MaxBackward1>)
tensor(0., device='cuda:0', grad_fn=<MinBackward1>)
tensor(0.1627, device='cuda:0', grad_fn=<MeanBackward0>)

ONNX

ONNX Out0 Max,Min,Mean:
1.0
0.44329888
0.9890943
ONNX Out1 Max,Min,Mean:
0.5567011
3.0826632e-14
0.010905608

Any ideas on how to get it working gratefully recieved!

[Advait Kawale (Advait251206)]

Hi,

One thing stands out in your code: you never put the PyTorch model into evaluation mode before running inference or exporting it.

Try calling eval() before both the PyTorch inference and ONNX export. This matters if the UNet contains layers such as Dropout or BatchNorm, whose behavior can differ between training and evaluation modes.

model_gc = torch.load("mymodel.pt", map_location=device)
model_gc.to(device)
model_gc.eval()

pic = pic.to(device)

with torch.inference_mode():
    pytorch_out = model_gc(pic)

torch.onnx.export(
    model_gc,
    pic,
    "mymodel.onnx",
    opset_version=11,
    input_names=["image"],
    output_names=["output"],
)

Then compare the outputs using the same input tensor:

with torch.inference_mode():
    pytorch_out = model_gc(pic).cpu().numpy()

ort_out = session.run(["output"], {"image": pic.cpu().numpy()})[0]

print("max abs difference:", np.max(np.abs(pytorch_out - ort_out)))
print("mean abs difference:", np.mean(np.abs(pytorch_out - ort_out)))

If the difference remains, I would check these next:

• Confirm preprocessing is identical. Your shown code passes the same pic tensor to both paths, which is good; make sure the saved model does not expect additional normalization or a different channel order.
• Compare the full tensors, not just min/max/mean. Those statistics can reveal a difference but do not show whether the outputs differ at every pixel or only in particular regions.
• Check the model’s output semantics. Your two channels appear complementary in both results, which may indicate a two-class softmax output, but that alone does not establish that the outputs are equivalent.
• If necessary, compare intermediate layer outputs to identify the first point where the PyTorch and ONNX results diverge. You can also try a newer supported ONNX opset/runtime combination after establishing a baseline.

The first test I’d run is simply model_gc.eval() followed by a fresh export and inference comparison. It is a common source of PyTorch-versus-ONNX discrepancies, especially for models containing training-sensitive layers.

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206

[dmagee]

Thanks for your answer. Since I posted this (~10 months ago) I worked out that the model actually only worked in train mode and not eval mode. Whoever trained it never tried it in eval mode, where there results were much different (inferior). It seems the model utilised Batch Normalisation layers as a kind of intensity/colour normalisation and so broke when put into eval mode. We ended up training a new model bootstrapped from the non-eval output of the original model on our data.