本篇记录一次从Pytorch图像处理模型转换成ONNX模型之后,在推理过程中出现了明显色偏问题的解决过程。
问题描述:原始pytorch模型推理正常,通过torch.onnx.export()函数转换成onnx之后,推理时出现了比较明显的颜色偏差。
原始模型转换程序如下:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
pth_path = 'model/my_model.pth'
onnx_path = 'model/my_model.onnx'
# 模型定义
model = MyModelStruct()
# 加载模型到设备
model.to(device)
# 加载checkpoint
checkpoint = torch.load(pth_path, map_location=device)
# 将checkpoint加载到模型
model.load_state_dict(checkpoint)
# 将模型设置为推理模式
model.eval()
# 定义模型输入输出
input_names = ['input', 'a', 'b']
output_names = ['output']
# 定义输入数据格式,随机数初始化
input = torch.rand(1,3, 512, 512)
a = torch.rand(1)
b = torch.rand(1)
# 将数据加载到设备
input = input.to(device)
a = a.to(device)
b = b.to(devici)
# 开始转换
torch.onnx.export(model, (input, a, b), onnx_path, input_names=input_names, output_names=output_names, verbose=True)
print('Done.')
1. 由于不是自己的模型,因此,这个pytorch模型拿到手后,先自己写了推理程序,在自己的PC上跑了一下。我自己的PC机只有CPU,在CPU上运行的结果跟onnx上一致,也存在色偏。但将同样的程序放到GPU服务器上运行,结果确实正常的。因此得出结论,应该是模型中间的某些算子,在CPU和GPU上的处理存在误差。
2. 将pytorch模型使用上面的转换程序转换成onnx之后,通过onnxruntime进行推理(这部分本篇先略过,后面再专门开一篇写onnxruntime推理),同样偏色,复现了最初的问题。
3. 通过研究torch.onnx.export()函数,发现跟算子处理关系最大的参数是opset_version,我的版本中,默认使用的opset_version为14,尝试换到16,转换出来的模型推理结果竟然正常了!
相应转换代码替换为:
torch.onnx.export(model, (input, a, b), onnx_path, input_names=input_names, output_names=output_names, verbose=True, opset_version=16)
torch.onnx.export()接收的部分关键参数解释如下:
_export(
model,
args,
f,
export_params,
verbose,
training,
input_names,
output_names,
operator_export_type=operator_export_type,
opset_version=opset_version,
do_constant_folding=do_constant_folding,
dynamic_axes=dynamic_axes,
keep_initializers_as_inputs=keep_initializers_as_inputs,
custom_opsets=custom_opsets,
export_modules_as_functions=export_modules_as_functions,
)
其中,
4. 到第三步其实还没完,因为,虽然模型输出没有偏色了,但是拿到实际场景中去运行,发现它的运行速度变慢了很多。于是又尝试了使用opset_version 11:
torch.onnx.export(model, (input, a, b), onnx_path, input_names=input_names, output_names=output_names, verbose=True, opset_version=11)
这次,也没有偏色,而且推理速度回归正常。
目前暂时使用opset_version 11,至于为什么版本16会使运行速度变慢,还需要更深入地去了解不同版本地差异。留待后续吧。