问题描述#xff1a;原始pytorch模型推理正常#xff0c;通过torch.onnx.export()函数转换成onnx之后#xff0c;推理时出现了比较明显的颜色偏差。

原始模型…本篇记录一次从Pytorch图像处理模型转换成ONNX模型之后在推理过程中出现了明显色偏问题的解决过程。

问题描述原始pytorch模型推理正常通过torch.onnx.export()函数转换成onnx之后推理时出现了比较明显的颜色偏差。

原始模型转换程序如下

model.load_state_dict(checkpoint)

将模型设置为推理模式

由于不是自己的模型因此这个pytorch模型拿到手后先自己写了推理程序在自己的PC上跑了一下。

我自己的PC机只有CPU在CPU上运行的结果跟onnx上一致也存在色偏。

但将同样的程序放到GPU服务器上运行结果确实正常的。

因此得出结论应该是模型中间的某些算子在CPU和GPU上的处理存在误差。

将pytorch模型使用上面的转换程序转换成onnx之后通过onnxruntime进行推理这部分本篇先略过后面再专门开一篇写onnxruntime推理同样偏色复现了最初的问题。

通过研究torch.onnx.export()函数发现跟算子处理关系最大的参数是opset_version我的版本中默认使用的opset_version为14尝试换到16转换出来的模型推理结果竟然正常了

相应转换代码替换为

torch.onnx.export()接收的部分关键参数解释如下

_export(model,args,f,export_params,verbose,training,input_names,output_names,operator_export_typeoperator_export_type,opset_versionopset_version,do_constant_foldingdo_constant_folding,dynamic_axesdynamic_axes,keep_initializers_as_inputskeep_initializers_as_inputs,custom_opsetscustom_opsets,export_modules_as_functionsexport_modules_as_functions,)

model需要被转换成onnx的模型。

args模型输入参数一般我们在这里指定模型输入数据的尺寸如果模型有多个参数该参数也可以是一个元组如本例中模型输入三个参数。

f要导出的onnx模型的路径包括onnx文件名。

export_params(bool,

default

True)为True时所有的模型参数都会被导出为False时则会导出一个未被训练的模型。

training(enum,

default

TrainingMode.EVAL)有三种模式分别为TrainingMode.EVALTrainingMOde.PRESERVETrainingMode.TRAINING一般使用EVAL模式即可。

input_names模型图中输入节点的名称字符串。

output_names模型图中输出节点的名称字符串。

operator_export_type算子导出类型包括OperatorExportTypes.ONNXOperatorExportTypes.ONNX_FALLTHROUGHOperatorExportTypes.ONNX_ATENOperatorExportTypes.ONNX_ATEN_FALLBACK我们常用的模型算子一般都有onnx支持因此默认选第一种。

opset_version这是我们本篇中问题解决的核心参数默认opset版本为14可用范围是7~16通过手动设置将其设为最高版本问题初步解决。

dynamic_axes允许在导出的onnx模型中创建变化的维度是一个字典形式默认为空。

到第三步其实还没完因为虽然模型输出没有偏色了但是拿到实际场景中去运行发现它的运行速度变慢了很多。

于是又尝试了使用opset_version

(input,

11至于为什么版本16会使运行速度变慢还需要更深入地去了解不同版本地差异。

留待后续吧。