详细步骤：pytorch pth转wts转tensorrt（自定义模型，不用parser）

模型转换与部署：从PyTorch到TensorRT的实战

最新推荐文章于 2025-09-17 08:47:10 发布

原创

最新推荐文章于 2025-09-17 08:47:10 发布 · 5k 阅读

46 ·

CC 4.0 BY-SA版权

文章标签：

#tensorrt #deep learning #cpp

这篇博客介绍了如何将PyTorch模型转换为权重文件(wts)，然后在资源有限的嵌入式设备上利用TensorRT进行模型优化和部署。作者详细阐述了两个步骤：首先，通过Python脚本将.pth模型转换为.wts格式，这个过程中涉及到权重的读取和格式转换；其次，讲解了如何在没有现成parser支持的情况下，使用TensorRT API手动构建网络并加载wts权重，最终生成TensorRT engine文件。整个过程对于在内存有限的设备上实现模型高效运行具有实际指导意义。

分两步转是考虑了如下应用场景：
需要把模型部署到内存有限的嵌入式板

自己的电脑上安装的有anaconda, pytorch等，但是在电脑上转的不能直接在板子上用
板子的内存有限，不能安装anaconda, pytorch这些，但是需要部署模型上去。
这时就可以现在电脑上把pth转成wts，再把wts传到板子上，在板子上转成tensorrt

(1) pth转wts
参考如下代码

import torch
from torch import nn
#load你的模型
import os
import struct

def main():
    net = torch.load('XXX.pth') #loadpth文件
    net = net.to('cuda:0')
    net.eval()
    
    f = open("XXX.wts", 'w') #自己命名wts文件
    f.write("{}\n".format(len(net.state_dict().keys())))  #保存所有keys的数量
    for k,v in net.state_dict().items():
        #print('key: ', k)
        #print('value: ', v.shape)
        vr = v.reshape(-1).cpu().numpy()
        f.write("{} {}".format(k, len(vr)))