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!