devbox
AllinToyou
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

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

作者：AllinToyou | 2024-04-05 18:10:05

赞

踩

将语义分割pytorch-onnx-om

在 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

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6

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])

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28

其中 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

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25

此时，说明我们的的 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")

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12

以下是对 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")

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71

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)

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45

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]))

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68

使用命令如下所示：

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

  • 1
  • 2
  • 3
  • 4
  • 5

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")

  • 1
  • 2
  • 3
  • 4

因为我们模型有 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

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16

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

  • 1
  • 2
  • 3
  • 4
  • 5

5.2.3 运行 ONNX 模型

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

  • 1

这里需要说明一下：

  • 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'")

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/AllinToyou/article/detail/367224
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号