模型部署——Pytorch模型转ONNX模型及测试转后ONNX模型（附代码）

一 、PyTorch 模型转换为 ONNX 模型注意事项

1.1 确保模型的输入和输出类型和形状一致

PyTorch 和 ONNX 具有不同的数据类型和格式，因此需要确保模型的输入和输出类型和形状在转换过程中保持一致。

可以使用 torch.onnx.export() 函数的 input_names 和 output_names 参数来指定模型的输入和输出名称。

可以使用 torch.onnx.export() 函数的 dynamic_axes 参数来指定模型的动态轴。

1.2 处理模型中的控制流

ONNX 不支持所有 PyTorch 的控制流操作，例如 if 语句和 while 循环。

需要对模型中的控制流进行改造，使其能够被 ONNX 理解和执行。

可以使用 torch.onnx.utils.convert_to_onnx() 函数来帮助转换控制流。

1.3 避免使用自定义操作符

ONNX 不支持所有 PyTorch 的自定义操作符。

如果模型中使用了自定义操作符，则需要将其转换为 ONNX 支持的操作符。

可以使用 torch.onnx.utils.convert_custom_ops() 函数来帮助转换自定义操作符。

1.4 检查模型的正确性

在将模型部署到生产环境之前，需要检查模型的正确性。

可以使用 onnx.checker.check_model() 函数来检查模型的结构和参数的有效性。

可以使用 onnxruntime 等工具来测试模型的推理性能。

二、安装必要包

在终端通过下面命令安装必要的onnx包和onnxruntime包：

pip install onnx -i https://pypi.mirrors.ustc.edu.cn/simple/

pip install onnxruntime -i https://pypi.mirrors.ustc.edu.cn/simple/

三、PyTorch 模型转换为 ONNX 模型

3.1 导入模型架构

将本教程代码拷贝到自己的项目工程文件中，找到模型架构位置并导入，如下：

3.2 代码参数修改

使用此教程代码需要修改的地方有：

3.3 代码

详细的转换代码见下：

import os.path
import onnx
import numpy as np
import onnxruntime
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.archs.ecbsr_arch import ECBSR
import torch.onnx

def pytorch_onnx_test(onnx_file_path: str, torch_model: torch.nn.Module):
    def to_numpy(tensor):  # 定义辅助函数
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    # 测试数据
    torch.manual_seed(66)  # 设置随机种子 ，确保每次运行产生相同随机输入数据。
    dummy_input = torch.randn(
        1, 3, 120, 90, device='cpu')  # 创建一个随机的 4D PyTorch 张量作为测试数据 张量形状为[1, 3, 120, 90]。将张量放到 CPU 上 (device='cpu')
    sess = onnxruntime.InferenceSession(onnx_file_path)  # 加载 ONNX 模型。      onnx_file_path为保存的 ONNX 模型文件路径
    # onnx 网络输出
    '''
    使用 sess.run(None, { "input": to_numpy(dummy_input)}) 运行 ONNX 模型。
    dummy_input 为前面创建的测试数据。
    将测试数据通过 to_numpy 函数转换为 NumPy 数组并作为输入传递给模型。
    None 表示不用指定模型的其他输入 (可能有多个输入)。
    该函数返回模型的输出，并将其转换为 NumPy 数组。
    '''
    onnx_out = np.array(sess.run(None, {"input": to_numpy(dummy_input)}))  #fc 输出是三维列表
    print("==============>")  # 打印分割线 ("==============>")
    print(onnx_out)  # 打印 ONNX 模型的输出 (onnx_out) 和其形状 (onnx_out.shape)
    print(onnx_out.shape)
    print("==============>")
    with torch.no_grad():  # 使用 with torch.no_grad(): 屏蔽梯度计算 (因为不需要计算梯度)。
        torch_model.eval()  # 将 PyTorch 模型 (torch_model) 设置为评估模式 (eval())
        torch_out_res = torch_model(dummy_input).detach().numpy(
        )  #fc输出是二维 列表    使用 torch_model(dummy_input) 运行 PyTorch 模型，得到输出 (torch_out_res)。分离 (detach) 并转换为 NumPy 数组
    print(torch_out_res)
    print(torch_out_res.shape)  # 打印 PyTorch 模型的输出 (torch_out_res) 和其形状 (torch_out_res.shape)。
    print("===================================>")
    print("输出结果验证小数点后四位是否正确,都变成一维np")
    torch_out_res = torch_out_res.flatten()  # 将 ONNX 和 PyTorch 模型的输出分别展平为一维数组 (flatten())。
    onnx_out = onnx_out.flatten()
    pytor = np.array(torch_out_res, dtype="float32")  #need to float32  将 PyTorch 和 ONNX 模型的输出转换为浮点型 (dtype="float32")。
    onn = np.array(onnx_out, dtype="float32")  ##need to float32
    decimal_arg = 4
    np.testing.assert_almost_equal(
        pytor, onn, decimal=decimal_arg
    )  #精确到小数点后4位，验证是否正确，不正确会自动打印信息.  使用 np.testing.assert_almost_equal 函数比较两个一维数组的元素，是否在小数点后 decimal_arg 位 (这里是 4 位) 相近,如果不相近，则会抛出异常并打印错误信息。
    print(
        f"恭喜你 ^^ , onnx 和 pytorch 结果一致, Exported model has been executed decimal={decimal_arg} and the result looks good!"
    )  # 否则，打印成功信息 ("恭喜你 ^^ , onnx 和 pytorch 结果一致...")。

## 设置预训练权重和 ONNX 模型保存路径
pretrained_weights = "experiments/pretrained_models/RealESRGAN_x4plus.pth"  # 定义预训练权重文件的路径,此文件包含训练好的模型参数。
# pretrained_weights = "experiments/RealECBSR_m4c32_L1loss/models/net_g_980000.pth"  # 定义预训练权重文件的路径,此文件包含训练好的模型参数。
onnx_save_path = os.path.join(
    "onnx_models", "RealESRGAN_RRDBx4_test.onnx")  # 定义导出的 ONNX 模型的保存路径，将模型保存为 RealESRGAN_RRDBx4.onnx 文件，位于 onnx_models 文件夹中。
if not os.path.exists(os.path.dirname(onnx_save_path)):
    os.makedirs(os.path.dirname(onnx_save_path))  #  检查 onnx_models 文件夹是否存在，如果不存在则创建该文件夹。
## 创建并加载 PyTorch 模型
torch_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32,scale=4)  # 创建一个名为 RRDBNet 的 PyTorch 模型实例。定义了模型的结构
# torch_model = ECBSR(num_in_ch=3, num_out_ch=3,  num_block=4, num_channel=32, num_grow_ch=32,scale=4,with_idt=True)
torch_model.load_state_dict(
    torch.load(pretrained_weights)['params_ema'],
    strict=True)  # 加载预训练权重文件。 从加载的字典中取出名为 'params_ema' 的键对应的值，该值包含模型的训练参数。strict=True: 严格模式，确保所有预训练权重的键都与模型架构兼容
torch_model.eval()  # 将模型设置为评估模式，以便稍后进行推理 (inference)
## 创建测试输入
batch_size = 1  # 定义批处理大小为 1，即一次只输入一张图片。
# Input to the model
x = torch.randn(
    batch_size, 3, 120, 90,
    requires_grad=True)  # 创建一个随机的 4D PyTorch 张量作为模型的输入。  requires_grad=True: 允许计算梯度 (稍后可能会用到，但在此处未使用)
torch_out = torch_model(x)
## 导出 PyTorch 模型为 ONNX 格式
# Export the model   使用 PyTorch 提供的 export 函数将 PyTorch 模型导出为 ONNX 格式
torch.onnx.export(
    torch_model,  # 要导出的 PyTorch 模型
    x,  # 模型的输入 (在此处使用创建的随机输入)
    onnx_save_path,  # 导出的 ONNX 模型的保存路径 (上面定义过)
    export_params=True,  # 将模型的参数嵌入到导出的 ONNX 模型中
    opset_version=11,  # 指定导出的 ONNX 模型兼容的 opset 版本 (运算符集版本)
    do_constant_folding=False,  # 禁用常量折叠优化 (需要根据实际情况调整)。
    input_names=['input'],  # 定义模型的输入名称为 "input"
    output_names=['output'],  # 定义模型的输出名称为 "output"
    dynamic_axes={
        'input': {
            0: 'batch_size'
        },  # 指定模型的动态轴
        'output': {
            0: 'batch_size'
        }
    })  # batch_size 可以是可变的，即可以一次输入多张图片
## 加载和检查导出的 ONNX 模型
onnx_model = onnx.load(onnx_save_path)  #  加载导出的 ONNX 模型
onnx.checker.check_model(onnx_model)  # 使用 onnx.checker.check_model 函数检查导出的 ONNX 模型是否有效
pytorch_onnx_test(onnx_save_path, torch_model)  #  调用 pytorch_onnx_test 函数测试导出的 ONNX 模型。onnx_save_path: 导出的 ONNX 模型的路径。

3.4 转换结果

3.5 运行输出

四、可视化模型

netron官网地址为：netron

通过netron可以直观的可视化onnx模型结构，利于理解与调整。打开上面网址的样子如下：

将自己的onnx模型直接拖到打开的网页上，如下：

五、总结

以上就是Pytorch模型转ONNX模型及测试转后ONNX模型的详细代码及过程，希望能帮到你。

总结不易，多多支持，谢谢！

感谢您阅读到最后！关注公众号「视觉研坊」，获取干货教程、实战案例、技术解答、行业资讯！