Pytorch图像模型转ONNX后出现色偏问题

本篇记录一次从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,
    )

其中，

• model：需要被转换成onnx的模型。
• args：模型输入参数，一般我们在这里指定模型输入数据的尺寸，如果模型有多个参数，该参数也可以是一个元组，如本例中模型输入三个参数。
• f：要导出的onnx模型的路径，包括onnx文件名。
• export_params：(bool, default True)，为True时所有的模型参数都会被导出；为False时，则会导出一个未被训练的模型。
• training：(enum, default TrainingMode.EVAL)，有三种模式，分别为：TrainingMode.EVAL，TrainingMOde.PRESERVE，TrainingMode.TRAINING，一般使用EVAL模式即可。
• input_names：模型图中输入节点的名称，字符串。
• output_names：模型图中输出节点的名称，字符串。
• operator_export_type：算子导出类型，包括：OperatorExportTypes.ONNX，OperatorExportTypes.ONNX_FALLTHROUGH，OperatorExportTypes.ONNX_ATEN，OperatorExportTypes.ONNX_ATEN_FALLBACK，我们常用的模型算子一般都有onnx支持，因此默认选第一种。
• opset_version：这是我们本篇中问题解决的核心参数，默认opset版本为14，可用范围是7~16，通过手动设置，将其设为最高版本，问题初步解决。
• dynamic_axes：允许在导出的onnx模型中创建变化的维度，是一个字典形式，默认为空。

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会使运行速度变慢，还需要更深入地去了解不同版本地差异。留待后续吧。