PyTorch语义分割模型转ONNX以及对比转换后的效果（PyTorch2ONNX、Torch2ONNX、pth2onnx、pt2onnx、修改名称、转换、测试、加载ONNX、运行ONNX）

在 PyTorch 官网已经给出了相关的文档，感兴趣的同学可以看一下文档：EXPORTING A MODEL FROM PYTORCH TO ONNX AND RUNNING IT USING ONNX RUNTIME

1. 准备工作

1. 语义分割模型 model.py
2. 训练好的权值文件 model.pth / model.pt
3. onnx==1.12.0
4. onnxruntime==1.15.1

import torch.onnx
from models import PPLiteSeg
import onnxruntime as ort
from PIL import Image
import numpy as np
import torchvision.transforms as transforms

2. 创建 PyTorch 模型

首先我们需要创建一个 PyTorch 模型并加载 .pth 权值文件：

# 创建模型
torch_model = PPLiteSeg()

# 加载模型权重
model_state_dict = torch.load("checkpoint/model.pth")

# 如果模型使用了DDP训练，则模型状态字典的会有'module'的前缀，我们需要删除
# 创建一个新的字典，去掉 "module." 前缀
# new_state_dict = {k.replace('module.', ''): v for k, v in model_state_dict['model'].items()}

# 加载模型权重
torch_model.load_state_dict(new_state_dict, strict=True)
print("\033[1;31m模型权重加载完毕...\033[0m")
    
"""
	因为我们的模型最终的输出并没有经过后处理，此时的shape为[N, num_classes, H, W]，所以需要对模型添加上后处理,
	让模型的输出为[N, 1, H, W]
"""
# 给模型添加后处理操作
torch_model = WrappedModel(torch_model)
    
# 设置模型为推理状态(这一步是必须的!)
torch_model.eval()
    
# 创建一个输入Tensor
x = torch.randn(1, 3, 512, 512, requires_grad=True)
torch_out = torch_model(x)
print(torch_out[0].shape)  # torch.Size([1, 1, 512, 512])

其中 WrappedModel 代码为：

import torch


class WrappedModel(torch.nn.Module):
    def __init__(self, model, output_op):
        super().__init__()
        self.model = model

    def forward(self, x):
        outs = self.model(x)
        new_outs = []
        for out in outs:
            out = torch.nn.functional.softmax(out, dim=1)  # 沿着通道维度进行概率计算
            label = torch.argmax(out, dim=1).to(dtype=torch.int32)  # 获取最大的位置
            label = torch.unsqueeze(label, 1)
            
            # torch.max返回值有两个：最大值的张量 + 最大值的索引张量
            max_score = torch.max(out, dim=1)[0]  # 获取最大概率
            max_score = torch.unsqueeze(max_score, 1)
            
            new_outs.append(label)
            new_outs.append(max_score)
		
		# 返回的是一个len==2的list
        return new_outs

此时，说明我们的的 PyTorch 模型创建成功并且正确地加载了训练好的权重。

3. 转换为 ONNX 模型并保存

# Export the model
torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "model.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
                  opset_version=11,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['label', 'score'], # the model's output names
                  dynamic_axes={'input' : {0 : 'B'},    # variable length axes
                                  'output' : {0 : 'B'}})
print("\033[1;31mONNX模型转换完毕.\033[0m")

以下是对 torch.onnx.export 函数的参数进行说明：

1. torch_model: 这是要导出的 PyTorch 模型的实例。

2. x: 这是模型的输入数据，可以是单个输入 Tensor 或一个包含多个输入 Tensor 的元组，取决于模型的输入方式。

3. "model.onnx": 这是导出的 ONNX 模型文件的保存路径。ONNX 模型将被保存在名为 “model.onnx” 的文件中。可以更改文件名和路径。

4. export_params=True: 这是一个布尔值，指示是否导出模型的参数权重。如果设置为 True，模型的参数将与模型一起保存到 ONNX 文件中，以便在推理时使用。如果设置为 False，则不会导出参数，仅导出模型结构。

5. opset_version=11: 这是导出模型所使用的 ONNX 版本。在此示例中，使用 ONNX 版本 11。不同版本的 ONNX 支持不同的操作，因此需要选择与的模型和运行时兼容的版本。

6. do_constant_folding=True: 这是一个布尔值，指示是否执行常量折叠以进行优化。如果设置为 True，则 ONNX 导出将尝试将模型中的常量 Tensor 折叠为常量节点，以减小模型文件的大小和提高推理速度。

7. input_names: 这是模型的输入名称列表（list），用于标识模型的输入 Tensor 。在此示例中，模型的输入 Tensor 被命名为 “input”。

8. output_names: 这是模型的输出名称列表（list），用于标识模型的输出 Tensor 。在此示例中，模型的输出 Tensor 被命名为 “label” 和 “score”。

9. dynamic_axes: 这是一个字典，用于指定动态轴的名称。动态轴是指可以具有可变长度的轴，通常是批处理轴。在此示例中，输入 “input” 和输出 “output” 的第一个维度被指定为 “B”，表示批处理轴可以具有可变长度。

通过使用这些参数，可以控制如何导出 PyTorch 模型到 ONNX 格式，并根据的需求进行配置。

说明：

1. 因为我们的模型的输出是一个长度为 2 的 list，所以 output_names 应该有两个；
2. dynamic_axes 表示哪些是动态的，这里我们将 Batch 维度设置为动态，即 ONNX 模型的 Batch 维度的输入是任意的，并非固定死的。

---

完整代码如下：

import torch
import numpy as np
import torch.onnx
from models import PPLiteSeg


class WrappedModel(torch.nn.Module):
    def __init__(self, model, output_op):
        super().__init__()
        self.model = model

    def forward(self, x):
        outs = self.model(x)
        new_outs = []
        for out in outs:
            out = torch.nn.functional.softmax(out, dim=1)  # 沿着通道维度进行概率计算
            label = torch.argmax(out, dim=1).to(dtype=torch.int32)  # 获取最大的位置
            label = torch.unsqueeze(label, 1)
            
            # torch.max返回值有两个：最大值的张量 + 最大值的索引张量
            max_score = torch.max(out, dim=1)[0]  # 获取最大概率
            max_score = torch.unsqueeze(max_score, 1)
            
            new_outs.append(label)
            new_outs.append(max_score)
		
		# 返回的是一个len==2的list
        return new_outs


if __name__ == "__main__":
	# 创建模型
	torch_model = PPLiteSeg()
	
	# 加载模型权重
	model_state_dict = torch.load("checkpoint/model.pth")
	
	# 如果模型使用了DDP训练，则模型状态字典的会有'module'的前缀，我们需要删除
	# 创建一个新的字典，去掉 "module." 前缀
	# new_state_dict = {k.replace('module.', ''): v for k, v in model_state_dict['model'].items()}
	
	# 加载模型权重
	torch_model.load_state_dict(new_state_dict, strict=True)
	print("\033[1;31m模型权重加载完毕...\033[0m")
	    
	"""
		因为我们的模型最终的输出并没有经过后处理，此时的shape为[N, num_classes, H, W]，所以需要对模型添加上后处理,
		让模型的输出为[N, 1, H, W]
	"""
	# 给模型添加后处理操作
	torch_model = WrappedModel(torch_model)
	    
	# 设置模型为推理状态(这一步是必须的!)
	torch_model.eval()
	    
	# 创建一个输入Tensor
	x = torch.randn(1, 3, 512, 512, requires_grad=True)
	torch_out = torch_model(x)
	
	# Export the model
	torch.onnx.export(torch_model,               # model being run
	                  x,                         # model input (or a tuple for multiple inputs)
	                  "model.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
	                  opset_version=11,          # the ONNX version to export the model to
	                  do_constant_folding=True,  # whether to execute constant folding for optimization
	                  input_names = ['input'],   # the model's input names
	                  output_names = ['label', 'score'], # the model's output names
	                  dynamic_axes={'input' : {0 : 'B'},    # variable length axes
	                                  'output' : {0 : 'B'}})
    print("\033[1;31mONNX模型转换完毕.\033[0m")

4. 修改 ONNX

4.1 修改输入输出的 shape

当我们保存为 ONNX 之后，我们可以使用一款名为 Netron 的软件打开 .onnx 文件，如下所示：

我们可以看到，ONNX 文件中的 ArgMax 对应的输出是 label，ReduceMax 对应的输出是 score，说明我们的模型转换是正确的。但是我们看右边会发现，input 的 shape 为 [B, 3, 512, 512]，这是我们想要的，但是输出按道理来说应该也是 [B, 3, 512, 512]，但并不是这样的。为了方便后期转换为 TRT（TensorRT），我们将输出进行修改，修改代码如下：

import onnx
import argparse


def show_inp_and_oup_info(model, modify=False):
    input_info = model.graph.input
    print("模型的输入信息：")
    for info in input_info:
        print(info.name, info.type)

    output_info = model.graph.output
    print("模型的输出信息：")
    for info in output_info:
        print(info.name, info.type)
        
    
if __name__ == "__main__":
    # 输入 ONNX 模型路径
    model_path = "model.onnx"

    # 输出 ONNX 模型路径
    output_path = "retype_model.onnx"

    # 读取 ONNX 模型
    model = onnx.load(model_path)
    
    show_inp_and_oup_info(model, modify=False)

    # 找到输入张量并修改
    # for input_info in model.graph.input:
    #     if input_info.name in ['x', 'input']:
    #         # 修改输入张量的形状
    #         input_info.type.tensor_type.shape.dim[0].dim_param = "B"

    # 修改输出张量的形状
    for output_info in model.graph.output:
        if output_info.name in ["label", "score"]:
            output_info.type.tensor_type.shape.dim[0].dim_param = "B"
            output_info.type.tensor_type.shape.dim[2].dim_value = 512
            output_info.type.tensor_type.shape.dim[3].dim_value = 512
            
    show_inp_and_oup_info(model, modify=True)
    
    # 保存修改后的模型
    onnx.save(model, output_path)

4.2 修改名称

如果我们想要对输入输出的名字进行修改，也可以使用下面的脚本：

import argparse
import sys
import onnx


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True, help='Path of directory saved the input model.')
    parser.add_argument('--origin_names', required=True, nargs='+', help='The original name you want to modify.')
    parser.add_argument('--new_names', required=True, nargs='+', 
                        help='The new name you want change to, the number of new_names should be same with the number of origin_names')
    parser.add_argument('--save_file', required=True, help='Path to save the new onnx model.')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()
    model = onnx.load(args.model)
    output_tensor_names = set()
    for ipt in model.graph.input:
        output_tensor_names.add(ipt.name)
    for node in model.graph.node:
        for out in node.output:
            output_tensor_names.add(out)

    for origin_name in args.origin_names:
        if origin_name not in output_tensor_names:
            print("[ERROR] Cannot find tensor name '{}' in onnx model graph.".format(origin_name))
            sys.exit(-1)
    if len(set(args.origin_names)) < len(args.origin_names):
        print("[ERROR] There's dumplicate name in --origin_names, which is not allowed.")
        sys.exit(-1)
    if len(args.new_names) != len(args.origin_names):
        print("[ERROR] Number of --new_names must be same with the number of --origin_names.")
        sys.exit(-1)
    if len(set(args.new_names)) < len(args.new_names):
        print("[ERROR] There's dumplicate name in --new_names, which is not allowed.")
        sys.exit(-1)
    for new_name in args.new_names:
        if new_name in output_tensor_names:
            print("[ERROR] The defined new_name '{}' is already exist in the onnx model, which is not allowed.")
            sys.exit(-1)

    for i, ipt in enumerate(model.graph.input):
        if ipt.name in args.origin_names:
            idx = args.origin_names.index(ipt.name)
            model.graph.input[i].name = args.new_names[idx]

    for i, node in enumerate(model.graph.node):
        for j, ipt in enumerate(node.input):
            if ipt in args.origin_names:
                idx = args.origin_names.index(ipt)
                model.graph.node[i].input[j] = args.new_names[idx]
        for j, out in enumerate(node.output):
            if out in args.origin_names:
                idx = args.origin_names.index(out)
                model.graph.node[i].output[j] = args.new_names[idx]

    for i, out in enumerate(model.graph.output):
        if out.name in args.origin_names:
            idx = args.origin_names.index(out.name)
            model.graph.output[i].name = args.new_names[idx]

    onnx.checker.check_model(model)
    onnx.save(model, args.save_file)
    print("[Finished] The new model saved in {}.".format(args.save_file))
    print("[DEBUG INFO] The inputs of new model: {}".format([x.name for x in model.graph.input]))
    print("[DEBUG INFO] The outputs of new model: {}".format([x.name for x in model.graph.output]))

使用命令如下所示：

python rename_onnx_model_name.py \
	   --model model.onnx \
	   --origin_names x y z \
	   --new_names x1 y1 z1 \
	   --save_file new_model.onnx

5. 测试转换前后效果

测试转换前后效果有两种思路：

1. 思路1：对比两种模型输出的差异 —— 机器看
2. 思路2：直接将两种模型的输出转换为图片 —— 肉眼看

5.1 对比两种模型输出的差异

在 PyTorch 教程中，是使用的这种方法。

# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(torch_res[0].numpy(), onnx_res[0], rtol=1e-03, atol=1e-05)
np.testing.assert_allclose(torch_res[1].detach().numpy(), onnx_res[1], rtol=1e-03, atol=1e-05)
print("\033[1;44mExported model has been tested with ONNXRuntime, and the result looks good!\033[0m")

因为我们模型有 score 和 label，所以两个都需要测试一下。

5.2 直接将两种模型的输出转换为图片

下面不进行具体的展示，只提供必要的函数。

5.2.1 加载图片并进行预处理

def load_test_img(image_path, target_size=(512, 512)):
    # 加载图片
    image = Image.open(image_path)

    # 调整图片大小为目标大小
    image = image.resize(target_size, Image.BILINEAR)

    # 使用 torchvision.transforms 将 PIL 图片转换为 PyTorch 张量
    transform = transforms.Compose([transforms.ToTensor(),  # 转换为张量
                                    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
    ])
    
    # 应用变换并添加批次维度 [1, C, H, W]
    image_tensor = transform(image).unsqueeze(0)

    return image_tensor

5.2.2 加载 ONNX 模型

def create_onnx_model(ckpt_path):
	import onnxruntime as ort
    ort_session = ort.InferenceSession(ckpt_path)
    print("\033[1;31mONNX模型创建完毕...\033[0m")
    return ort_session

5.2.3 运行 ONNX 模型

onnx_res = onnx_model.run(None, {"input": [test_img.squeeze(0)]})

这里需要说明一下：

• onnx_model.run: 这是运行 ONNX 模型的方法。onnx_model 是通过 ONNX Runtime 创建的 ONNX 模型的实例。

• None: 这是用于指定期望的输出名称的占位符。在此示例中，None 表示我们不指定输出名称，因此 ONNX Runtime 将返回所有输出。

• {"input": [test_img.squeeze(0)]}: 这是输入数据的字典。ONNX 模型通常需要一个字典来指定输入数据，其中键是输入名称，值是输入数据。在这里，输入名称为 “input”，对应的输入数据是 test_img.squeeze(0)。

test_img.squeeze(0): 这是将 test_img Tensor 的第一个维度（通常是批处理维度）挤压（去除），以便它符合 ONNX 模型的输入要求。通常，ONNX 模型的输入Tensor 期望没有批处理（Batch）维度，因此我们使用 .squeeze(0) 来去除第一个维度，以使输入数据与 ONNX 模型兼容。

运行此命令后，onnx_res 将包含 ONNX 模型的输出结果。这个结果通常是一个包含输出 Tensor 的列表（记住，是一个 list），其中每个元素对应一个模型输出。可以根据模型的输出情况来访问和处理这些结果。在这个特定示例中，可能需要进一步处理 onnx_res，以便将其转换为可用的数据或进行其他后续操作，具体取决于的应用场景。

5.2.4 将模型结果保存为图片

def save_torch_res(torch_res, suffix):
    # 转换 PyTorch 张量为 NumPy 数组
    torch_res_numpy = torch_res[0].squeeze(0).numpy()
    # 如果形状不是 [H, W]，可以进一步调整
    print(np.shape(torch_res_numpy))
    
    # 如果形状不是 [H, W]，可以进一步调整
    if torch_res_numpy.shape[0] == 1:
        torch_res_numpy = torch_res_numpy[0]

    # 创建灰度图像
    gray_image = Image.fromarray((torch_res_numpy * 255).astype('uint8'), mode='L')

    # 将灰度图像转换为伪彩色图像（伪彩色映射可根据需要更改）
    pseudo_color_image = gray_image.convert('P', palette=Image.ADAPTIVE, colors=256)

    # 保存伪彩色图像
    pseudo_color_image.save("results/pytorch_output_pseudo_color_image.png")
    print("伪彩色图像已保存为 'results/pytorch_output_pseudo_color_image.png'")


def save_onnx_res(onnx_res, suffix):
    # 转换 ONNX 结果为 NumPy 数组
    onnx_res_numpy = np.array(onnx_res[0])
    

    # 如果形状不是 [H, W]，可以进一步调整
    if onnx_res_numpy.shape[0] == 1:
        onnx_res_numpy = np.squeeze(onnx_res_numpy, axis=0)
        onnx_res_numpy = np.squeeze(onnx_res_numpy, axis=0)
    
    # 创建灰度图像
    gray_image = Image.fromarray((onnx_res_numpy * 255).astype('uint8'), mode='L')

    # 将灰度图像转换为伪彩色图像（伪彩色映射可根据需要更改）
    pseudo_color_image = gray_image.convert('P', palette=Image.ADAPTIVE, colors=256)

    # 保存伪彩色图像
    pseudo_color_image.save("results/onnx_output_pseudo_color_image.png")
    print("伪彩色图像已保存为 'results/onnx_output_pseudo_color_image.png'")