TensorRT学习（1）：通过pth生成wts文件

TensorRT学习（1）：通过pth生成wts文件

1. pth文件简介

pth文件是pytorch保存模型的一种方式，该文件只保存模型的参数。模型参数实际上一个字典类型，通过key-value的形式存储。

2. wts文件格式示例（来自wang-xinyu大佬的github）

第一行代表该文件有多少行，不包括它本身。之后每行格式为：

  [weight name] [value count = N] [value1] [value2], ..., [valueN]

10		
conv1.weight 150 be40ee1b bd20bab8 bdc4bc53 .......
conv1.bias 6 bd327058 .......
conv2.weight 2400 3c6f2220 3c693090 ......
conv2.bias 16 bd183967 bcb1ac8a .......
fc1.weight 48000 3c162c20 bd25196a ......
fc1.bias 120 3d3c3d49 bc64b948 ......
fc2.weight 10080 bce095a4 3d33b9dc ......
fc2.bias 84 bc71eaa0 3d9b276c ....... 
fc3.weight 840 3c252870 3d855351 .......
fc3.bias 10 bdbe4bb8 3b119ee0 ......

3. pth转wts

完整代码如下：

 import torch
 import struct
 
 net = torch.load('pth_path')	# 加载pth文件，pth_path为pth文件的路径、
 
 #for k in net.keys():			# 打印net的key
 #   print(k)		# 我的为epoch、instance_acc、class_acc、model_state_dict、optimizer_state_dict

 # 显然，模型参数保存在 net["model_state_dict"]，有时候命名会不一样，所以应打印net的key确定一下
 model_state_dict = net["model_state_dict"]
 f = open("pointNet_sem.wts", 'w')  # 自己命名wts文件
 f.write("{}\n".format(len(model_state_dict.keys())))  # 保存所有keys的数量
    for k, v in model_state_dict.items():
        vr = v.reshape(-1).cpu().numpy()	# 权重参数展开成1维
        f.write("{} {}".format(k, len(vr)))  # 保存每一层名称和参数长度
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())  # 使用struct把权重封装成字符串
        f.write("\n")