Hi there, I tried to export a small pretrained (fashion MNIST) model to ONNX for test cases and evaluated the results. The outputs were completely differnt and I already tried different solutions which did not help to solve the problem.
I spent weeks on finding a solution to this error - may you please help me?
Here are the labels and a formatted output of the PyTorch model and the ONNX model which was run with ONNX Runtime
Labels: ['T-shirt', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', ' Bag', 'Ankle boot']
PyTorch output: [14.139054 -1.463651 3.0200443 5.6304793 -3.21046 -6.91372 10.828293 -12.109244 -0.4888268 -7.0172133]
ONNX Runtime Output: [5.685347 4.6847196 -4.7160497 -11.215544 -14.1343975 28.269392 1.3876779 -3.9011323 1.518218 0.49619618]
Below you can find the unformatted output and the used files.
Unformatted output
Label: T-shirt Confidence: 0.9645941 Time taken: 0.07566659999999992
PyTorch output: [ 14.139054 -1.463651 3.0200443 5.6304793 -3.21046 -6.91372
10.828293 -12.109244 -0.4888268 -7.0172133]
ONNX Runtime output: [ 5.685347 4.6847196 -4.7160497 -11.215544 -14.1343975
28.269392 1.3876779 -3.9011323 1.518218 0.49619618]
========================================
Final top prediction is: Sandal with confidence of: 1.0
========================================
========================================
Inference time: 0.99 ms
========================================
Export routine
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
RAND_INPUT = torch.randn(BATCH_SIZE, 3, 112, 112, device=DEVICE)
torch.onnx.export(MODEL, # model being run
RAND_INPUT, # model input (or a tuple for multiple inputs)
str(MODEL.__class__.__name__)+".onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
#verbose=True, # output graph to console
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
enable_onnx_checker=True,
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes
'output' : {0 : 'batch_size'}})
Neural Network Model (mnist_model.py)
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import numpy as np
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return x
Testing routine (test.py)
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
def bulkPredictImages(model: nn.Module, classes: list, dataLoader: DataLoader):
resultArray = []
resultConfidenceArray = []
totalTime = 0.0
for images, _ in dataLoader:
if torch.cuda.is_available():
images = images.cuda()
for image in images:
image = image.unsqueeze_(0)
result, confidence, timeTaken = predictImage(model, classes, image)
resultArray.append(result)
resultConfidenceArray.append(confidence)
totalTime += timeTaken
return resultArray, resultConfidenceArray, totalTime
def predictImage(model: nn.Module, classes: list, inputImage: Image, transform = None):
timeStart = time.perf_counter() #in seconds
model.eval()
with torch.no_grad():
if transform != None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inputImage = transform(inputImage).float()
inputImage = inputImage.unsqueeze_(0)
inputImage = inputImage.to(device)
output = model(inputImage)
idx = output.data.cpu().numpy().argmax()
label = classes[idx]
top_p, top_classes = F.softmax(output, dim=1).topk(1, dim = 1)
confidence = top_p.data.cpu().numpy().max()
timeStop = time.perf_counter() #in seconds
return label, confidence, timeStop - timeStart
def main():
pass
if __name__ == '__main__':
main()
Converting and evaluation (PyTorchToOnnxConverter.py) (please have mercy for my coding style)
import io
import os
import sys
import numpy as np
import time
from torch import nn
import torch.onnx
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torchvision import transforms
from collections import OrderedDict
import onnx
import onnxruntime
from onnx import optimizer as optim
from PIL import Image
import test
import mnist_model
MODEL = None
BATCH_SIZE = 1
MODEL_PATH = ""
TEST_IMAGE_PATH = ""
LABELS = []
IMG = None
MONOCHROME = False
INPUT_LAYOUT = ()
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
RAND_INPUT = None
TRANSFORM = None
def setMNIST():
global MODEL
global MODEL_PATH
global TEST_IMAGE_PATH
global LABELS
global IMG
global MONOCHROME
global INPUT_LAYOUT
global RAND_INPUT
global TRANSFORM
MODEL = mnist_model.Net()
if torch.cuda.is_available():
MODEL = MODEL.cuda()
MODEL_PATH = "D:/Projekte/Masterthesis/DeepNerualNetworks/CustomFashionMNIST/best_model.pth"
TEST_IMAGE_PATH = "D:/Projekte/Masterthesis/DeepNerualNetworks/PyTorchToOnnxConverter/tshirt.png"
LABELS = ['T-shirt', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', ' Bag', 'Ankle boot']
IMG = Image.open(TEST_IMAGE_PATH).convert('L')
INPUT_LAYOUT = (1, 1, 28, 28)
MONOCHROME = True
RAND_INPUT = torch.randn(BATCH_SIZE, 1, 28, 28, device=DEVICE)
TRANSFORM = transforms.Compose([
transforms.Resize(size=28),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
def preprocess(input_data):
# convert the input data into the float32 input
img_data = input_data.astype('float32')
#normalize
mean_vec = np.array([0.485, 0.456, 0.406])
stddev_vec = np.array([0.229, 0.224, 0.225])
norm_img_data = np.zeros(img_data.shape).astype('float32')
for i in range(img_data.shape[0]):
norm_img_data[i,:,:] = (img_data[i,:,:]/255 - mean_vec[i]) / stddev_vec[i]
#add batch channel
norm_img_data = norm_img_data[np.newaxis, :, :, :]
#norm_img_data = img_data[np.newaxis, :, :, :]
norm_img_data = np.resize(norm_img_data, INPUT_LAYOUT).astype('float32')
return norm_img_data
def softmax2(arr):
sum = 0
max = np.max(arr)
print(max)
for i in arr:
print(i)
sum += np.exp(i - max) # subtract max in numr and denr to avoid overflow
res = []
for i in arr:
res.append(np.exp(i - max) / sum)
return res
def softmax3(inp):
top_p, top_classes = F.softmax(inp, dim=1).topk(1, dim = 1)
confidence = top_p.data.cpu().numpy().max()
return top_p, confidence
def softmax(x):
sum = 0
print("ONNX Runtime output: " + str(x[0][0]))
x = x.reshape(-1)
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum(axis=0)
def postprocess(result):
smRes = softmax(np.array(result)).tolist()
return smRes
def onnxImageProcess(img, monochrome: bool = False):
image_array = np.asarray(img)
if monochrome:
image_array = image_array[:, :, np.newaxis]
image_array = image_array.transpose(2, 0, 1)
return preprocess(image_array)
def runOnnxTest(ort_session, input_data):
#image = Image.open(TEST_IMAGE_PATH)
input_name = ort_session.get_inputs()[0].name
start = time.time()
raw_result = ort_session.run(None, {input_name: input_data})
end = time.time()
res = postprocess(raw_result)
#res = softmax2(raw_result)
#softmax3(raw_result)
inference_time = np.round((end - start) * 1000, 2)
idx = np.argmax(res)
print('========================================')
print('Final top prediction is: ' + LABELS[idx] + " with confidence of: " + str(res[idx]))
print('========================================')
print('========================================')
print('Inference time: ' + str(inference_time) + " ms")
print('========================================')
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
def torchImageProcess(img, transform):
inputImage = transform(img).float()
inputImage = inputImage.unsqueeze_(0)
return inputImage.to(DEVICE)
def compareOnnxModelWithTorch(session, onnxInput, torchOutput):
ort_inputs = {session.get_inputs()[0].name: onnxInput}
ort_outs = session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torchOutput), ort_outs[0], rtol=1e-03, atol=1e-05)
def add_value_info_for_constants(model : onnx.ModelProto):
"""
Currently onnx.shape_inference doesn't use the shape of initializers, so add
that info explicitly as ValueInfoProtos.
Mutates the model.
Args:
model: The ModelProto to update.
"""
# All (top-level) constants will have ValueInfos before IRv4 as they are all inputs
if model.ir_version < 4:
return
def add_const_value_infos_to_graph(graph : onnx.GraphProto):
inputs = {i.name for i in graph.input}
existing_info = {vi.name: vi for vi in graph.value_info}
for init in graph.initializer:
# Check it really is a constant, not an input
if init.name in inputs:
continue
# The details we want to add
elem_type = init.data_type
shape = init.dims
# Get existing or create new value info for this constant
vi = existing_info.get(init.name)
if vi is None:
vi = graph.value_info.add()
vi.name = init.name
# Even though it would be weird, we will not overwrite info even if it doesn't match
tt = vi.type.tensor_type
if tt.elem_type == onnx.TensorProto.UNDEFINED:
tt.elem_type = elem_type
if not tt.HasField("shape"):
# Ensure we set an empty list if the const is scalar (zero dims)
tt.shape.dim.extend([])
for dim in shape:
tt.shape.dim.add().dim_value = dim
# Handle subgraphs
for node in graph.node:
for attr in node.attribute:
# Ref attrs refer to other attrs, so we don't need to do anything
if attr.ref_attr_name != "":
continue
if attr.type == onnx.AttributeProto.GRAPH:
add_const_value_infos_to_graph(attr.g)
if attr.type == onnx.AttributeProto.GRAPHS:
for g in attr.graphs:
add_const_value_infos_to_graph(g)
return add_const_value_infos_to_graph(model.graph)
def remove_initializer_from_input(model):
if model.ir_version < 4:
print(
'Model with ir_version below 4 requires to include initilizer in graph input'
)
return
inputs = model.graph.input
name_to_input = {}
for input in inputs:
name_to_input[input.name] = input
for initializer in model.graph.initializer:
if initializer.name in name_to_input:
inputs.remove(name_to_input[initializer.name])
return model
def main():
setMNIST()
state = 0
so = onnxruntime.SessionOptions()
so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
state = torch.load(MODEL_PATH, map_location=DEVICE)
new_state_dict = OrderedDict()
for k, v in state['state_dict'].items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
MODEL.load_state_dict(new_state_dict)
MODEL.eval()
MODEL.is_training = False
#torch.save(MODEL, "512.pt")
result, confidence, timeTaken = test.predictImage(MODEL, LABELS, IMG, TRANSFORM)
sys.stdout.write("Label: " + str(result) + " Confidence: " + str(confidence) + " Time taken: " + str(timeTaken) + "\n")
#x = torch.randn(BATCH_SIZE, 3, 112, 112, requires_grad=True)
# Export the model
torch.onnx.export(MODEL, # model being run
RAND_INPUT, # model input (or a tuple for multiple inputs)
str(MODEL.__class__.__name__)+".onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
#verbose=True, # output graph to console
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
enable_onnx_checker=True,
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes
'output' : {0 : 'batch_size'}})
onnx_model = onnx.load(str(MODEL.__class__.__name__)+".onnx")
onnx_model = remove_initializer_from_input(onnx_model)
#passes = ["extract_constant_to_initializer", "eliminate_unused_initializer"]
#add_value_info_for_constants(onnx_model)
#for init in onnx_model.graph.initializer:
# for value_info in onnx_model.graph.value_info:
# if init.name == value_info.name:
# onnx_model.graph.input.append(value_info)
#onnx_model = optim.optimize(onnx_model, passes)
onnx.save(onnx_model, str(MODEL.__class__.__name__)+".onnx")
# Check that the IR is well formed
onnx.checker.check_model(onnx_model)
# Print a human readable representation of the graph
graph = onnx.helper.printable_graph(onnx_model.graph)
with open("graph_output.txt", mode="w") as fout:
fout.write(graph)
ort_session = onnxruntime.InferenceSession(str(MODEL.__class__.__name__)+".onnx", so)
compareOnnxModelWithTorch(ort_session, to_numpy(RAND_INPUT), MODEL(RAND_INPUT))
#compareOnnxModelWithTorch(ort_session, onnxImageProcess(IMG, MONOCHROME), MODEL(torchImageProcess(IMG, TRANSFORM)))
res = to_numpy(MODEL(torchImageProcess(IMG, TRANSFORM)))
print("PyTorch output: " + str(res[0]))
runOnnxTest(ort_session, onnxImageProcess(IMG, MONOCHROME))
if __name__ == "__main__":
main()
Thank you for your time and help
