[人工智能]详细步骤：pytorch pth转wts转tensorrt（自定义模型，不用parser）

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

• 自己的电脑上安装的有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)))  #保存每一层名称和参数长度
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())  #使用struct把权重封装成字符串
        f.write("\n")

if __name__ == '__main__':
    main()

wts文件是如下格式
第一行：数字
下面分别是网络每一层的名称，参数个数和对应的参数

(2) wts转tensorrt
常用的模型转换可参考链接
但是如果是自定义模型，不在常用模型范围，比如我建了一个几层的小语义分割网络，但是现有的parser不支持，那只能手撕tensorrt的API了

tensorRT的API使用方法可参考链接

简要说下wts转tensorrt的原理

1. 从wts文件把weight给load出来，存到一个map里，key是网络每层的名称，value就是对应的权重
2. 利用tensorrt的API把网络重建出来，同时导入key对应的value，也就是weightMap的形式
3. 定义网络的输出，设置内存空间
4. build engine
   最后出来的是一个engine文件

具体的做法：

1. load weight

std::map<std::string, Weights> loadWeights(const std::string file)
{
    std::cout << "Loading weights: " << file << std::endl;
    std::map<std::string, Weights> weightMap;

    // Open weights file
    std::ifstream input(file);
    assert(input.is_open() && "Unable to load weight file.");

    // Read number of weight blobs
    int32_t count;
    input >> count;
    assert(count > 0 && "Invalid weight map file.");

    while (count--)
    {
        Weights wt{DataType::kFLOAT, nullptr, 0};
        uint32_t size;

        // Read name and type of blob
        std::string name;
        input >> name >> std::dec >> size;
        wt.type = DataType::kFLOAT;

        // Load blob
        uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
        for (uint32_t x = 0, y = size; x < y; ++x)
        {
            input >> std::hex >> val[x];
        }
        wt.values = val;
        
        wt.count = size;
        weightMap[name] = wt;
    }

    return weightMap;
}

2. 重建网络
   具体参考tensorRT的API
   举个例子：
   比如是3x3的conv，stride=2, padding=2, kernel size=5x5，output channel=8
   首先要定义一个network

INetworkDefinition* network = builder->createNetworkV2(0U); 

输入层要定义出来，比如输入size是3xHxW

ITensor* data = network->addInput(INPUT_BLOB_NAME, DataType::kFLOAT, Dims3{3, INPUT_H, INPUT_W}); 

然后把输入传给conv层

IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 8, DimsHW{5, 5}, weightMap["conv1.weight"], weightMap["conv1.bias"]); //map名称要改

一层一层传下去，直到指定output

3. 定义输出，指定内存空间

deconv3->getOutput(0)->setName(OUTPUT_BLOB_NAME); //设置output名称
network->markOutput(*deconv3->getOutput(0)); //指定output

还要设置一个engine内存空间，太小的话会报错，尤其是报could not find implementation for node的error

builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(128*(1 << 20)); 

4. build engine

ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

这里说明一下反卷积层的写法，torch的deconvolution有padding和output_padding，如何设置
padding对应的是setPrePadding
output_padding对应的是setPostPadding

举个例子，反卷积层，output channel=16，kernel_size=5x5,
stride=2, padding=2, output_padding=1

IDeconvolutionLayer* deconv1 = network->addDeconvolutionNd(*relu3->getOutput(0), 16, DimsHW{5,5}, weightMap["deconv1.weight"], weightMap["deconv1.bias"]);
deconv1->setStrideNd(DimsHW{2, 2});
deconv1->setPrePadding(DimsHW{2, 2});
deconv1->setPostPadding(DimsHW{1, 1});

具体cpp和CMakeList.txt放在github上