Pytroch保存的.model/.pkl等模型如何转成.onnx模型，并最终转换成TNN模型，并完成移植

Audio·Wang

已于 2023-11-30 11:33:28 修改

阅读量775

点赞数 7

文章标签： python

于 2023-11-28 11:13:12 首次发布

本文链接：https://blog.csdn.net/aaaasong/article/details/134658791

版权

开头：最近开始搞模型抑制了，最终目标是要把算法转换成C++，具体过程还在摸索中，大致了解到过程需要pytorch2onnx→onnx2tnn。。。先开个帖子记录一下，后续慢慢补全。

第一步：将Pytroch保存的训练好的.model/.pkl等模型转成.onnx模型

将 PyTorch 模型导出为 ONNX 模型，可以使用 PyTorch 提供的 torch.onnx.export() 函数。

import torch
from XXX import Model

#加载已经训练好的模型
model = Model
Load_Path = ''  # 训练好的模型路径
torch.load(Load_Path, map_location=torch.device('cpu'))
model.load_state_dict()
model.eval()  # 设置为评估模式

# 指定输入大小，这里以示例输入为准
input = torch.randn(1, 2, 224, 224)  # 输入大小为(batch_size, channels, t, f)

# 指定导出的文件路径
onnx_path = "model.onnx"

# 导出模型为 ONNX 格式
torch.onnx.export(model, input, onnx_path, verbose=True)

print(f"ONNX 模型已成功导出到 {onnx_path