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 对应的输出是 labelReduceMax 对应的输出是 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")

因为我们模型有 scorelabel,所以两个都需要测试一下。

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'")
  • 1
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 4
    评论
### 回答1: 要将PyTorch的.pth模型转换ONNX格式,可以使用以下步骤: 1. 安装ONNX包:在终端中运行`pip install onnx`来安装ONNX包。 2. 加载PyTorch模型:使用PyTorch加载模型并将其转换ONNX格式。例如,可以使用以下代码加载PyTorch模型: ``` import torch import torchvision # 加载PyTorch模型 model = torchvision.models.resnet18(pretrained=True) # 将模型转换为eval模式 model.eval() # 创建一个虚拟输入张量 input_tensor = torch.randn(1, 3, 224, 224) # 导出模型ONNX格式 torch.onnx.export(model, input_tensor, "resnet18.onnx") ``` 在这个例子中,我们加载了一个预训练的ResNet18模型,并将其转换ONNX格式。我们还创建了一个虚拟输入张量,它将用于导出模型。 3. 验证ONNX模型:可以使用ONNX Runtime或其他ONNX兼容的推理引擎来验证导出的ONNX模型。例如,可以使用以下代码验证导出的ResNet18模型: ``` import onnx import onnxruntime # 加载ONNX模型 onnx_model = onnx.load("resnet18.onnx") # 验证模型 onnx.checker.check_model(onnx_model) # 创建一个ONNX Runtime会话 ort_session = onnxruntime.InferenceSession("resnet18.onnx") # 运行模型 ort_inputs = {ort_session.get_inputs()[].name: input_tensor.numpy()} ort_outputs = ort_session.run(None, ort_inputs) # 打印输出 print(ort_outputs) ``` 在这个例子中,我们使用ONNX Runtime创建了一个会话,并使用虚拟输入张量运行模型。我们还打印了模型的输出。 ### 回答2: PyTorch是流行的深度学习框架之一,而ONNX(Open Neural Network Exchange)是一个开源的跨平台深度学习框架,可以方便地让用户在多个平台上运行模型。在机器学习和深度学习应用中,模型转换和共享非常重要,而pth模型ONNX模型可以更加高效地在多个平台上利用训练好的模型。本文将介绍如何使用PyTorchpth模型转换ONNX模型。 首先,需要安装PyTorchONNX。在安装好这两个框架之后,使用以下代码将pth模型转换ONNX模型: ``` import torch import onnx # 加载pth模型 model = torch.load('model.pth') # 将pth模型转换onnx模型 input_names = ['input'] output_names = ['output'] dynamic_axes = {'input':{0:'batch_size'},'output':{0:'batch_size'}} x = torch.randn(1, 3, 224, 224) torch.onnx.export(model, x, "model.onnx", input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes) ``` 在此代码片段中,`input_names`和`output_names`分别表示网络的输入和输出节点名称。`dynamic_axes`参数确定哪些轴是变化的,这是非常重要的,因为不同的框架可能需要特定的推理数据格式。在这个例子中,`dynamic_axes`参数将输入和输出数据的第一维指定为“batch_size”,因为第一维是数据样本数量,可能因推断过程而变化。 在代码运行完毕后,可以得到一个ONNX模型,可以将其部署到ONNX支持的设备上进行推理,例如移动设备和边缘计算设备。 需要注意的是,将pth模型转换ONNX模型有一些限制和注意事项: 1. PyTorch支持的操作不一定是ONNX支持的。要将模型成功转换ONNX格式,使用的PyTorch操作必须是ONNX支持的。 2. ONNX不支持大部分运行时计算。因此,如果使用了一些需要计算图中的其他参数的操作(如动态图),则不能将模型成功转换ONNX格式。 3. 一些PyTorch操作(如Dropout)在训练和推断过程中有不同的行为。因此,需要在代码中明确指定模型的模式,以确保在转换和推理过程中得到正确的结果。 综上所述,pth模型ONNX模型可以方便地在多个平台上部署和使用训练好的模型,需要仔细注意输入输出节点、动态轴和框架支持的操作等细节。 ### 回答3: PyTorch是一个非常流行和广泛使用的深度学习框架。在深度学习中,常常需要将训练得到的模型转换为一种可移植的格式,如ONNX。这样可以让我们在不同的框架和设备上使用模型。 将PyTorch模型转换ONNX需要以下步骤。 第一步是安装必要的库和工具包。我们需要安装最新版本的PyTorchONNX。可以在PyTorch的官方网站和ONNX的GitHub页面上找到安装说明。 第二步是准备PyTorch模型。在PyTorch中,我们可以使用torch.save将模型保存为.pth文件。这个文件包含了模型的权重和架构。 第三步是使用torch.onnx.export将模型转换ONNX格式。这个函数需要指定PyTorch模型、输入张量和输出文件的路径。我们还可以使用其他参数来设置转换的选项,如输入和输出名称、数据类型和设备类型等。 第四步是验证转换的结果。我们可以使用ONNX Runtime或其他支持ONNX格式的框架加载模型,输入数据并进行推理。通过比较转换前后的输出,我们可以确认转换的正确性。 需要注意的是,PyTorch模型ONNX模型在一些细节上可能存在差异。例如,PyTorch中的一些操作可能没有对应的ONNX实现,或者ONNX中的一些操作可能需要特定的属性和参数。因此,在进行模型转换时,需要了解两种格式的差异,并根据实际需求进行调整。 总之,PyTorch模型转换ONNX格式可以让我们更加方便地在不同的框架和设备上使用模型,是深度学习工作中不可或缺的一部分。
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值