TensorRT基础使用

*Major*

已于 2022-10-10 11:07:46 修改

阅读量353

点赞数

文章标签： trt

于 2022-09-12 15:01:39 首次发布

本文链接：https://blog.csdn.net/qq_41375318/article/details/126817458

版权

TensorRT基础使用

版本8.2.3GA(成熟稳定版)和8.4.0EA（新功能测试版）以后，有C++和python的API，完全等价可以混用

官方教程：

https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#overview

https://docs.nvidia.com/deeplearning/tensorrt/api/c_api

https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/

https://github.com/NVIDIA/trt-samples-for-hackathon-cn/tree/master/cookbook

构建期(推理优化器)

1.模型解析/建立
加载 Onnx等其他格式的模型/使用原生API搭建模型
2.计算图优化
横向层融合(Conv)，纵向层融合(Conv+add+ReLU)，……
3.节点消除
去除无用层，节点变换（Pad,Slice，Concat，Shuffle),.……
4.多精度支持
FP32/ FP16/ INT8 / TF32(可能插入reformat节点)
5.优选 kernel / format
硬件有关优化
6.导入plugin
实现自定义操作
7.显存优化
显存池复用

运行期(运行时环境)

1.运行时环境
对象生命期管理，内存显存管理，异常处理
2.序列化反序列化
推理引擎保存为文件或从文件中加载

使用流程

    // 1.编译模型：onnx-->trtmodel
    TRT::compile(
        TRT::Mode::FP32,//模式，fp32,fp16,int8
        1, // 最大batchsize
        "test.onnx", // onnx文件，输入
        "test.fp32.trtmodel" // trt输出模型
    )

    // 2.加载编译好的引擎
    auto infer = TRT::load_infer("test.fp32.trtmodel");
       
    // 3.设置输入的值
    infer->input(0)->set_to(1.0f);

    // 4.引擎推理
    infer->forward();

    // 5.取出引擎的输出并打印
    auto out = infer->output(0);
    for (int i = 0;i<out->channel();++i)
    {
        INFO("%f",out->at<float>(0,i));
    }

#define _CRT_SECURE_NO_WARNINGS；
// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>

// cuda include
#include <cuda_runtime.h>

// system include
#include <stdio.h>

class TRTLogger : public nvinfer1::ILogger {


public:
    virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {
        if (severity <= Severity::kVERBOSE) {
            printf("%d: %s\n", severity, msg);
        }
    }
};

nvinfer1::Weights make_weights(float* ptr, int n) {
    nvinfer1::Weights w;
    w.count = n;
    w.type = nvinfer1::DataType::kFLOAT;
    w.values = ptr;
    return w;
}

int main() {
    // 本代码主要实现一个最简单的神经网络 figure/simple_fully_connected_net.png 

    TRTLogger logger; // logger是必要的，用来捕捉warning和info等

    // ----------------------------- 1. 定义 builder, config 和network -----------------------------
    // 这是基本需要的组件
    //形象的理解是你需要一个builder去build这个网络，网络自身有结构，这个结构可以有不同的配置
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    // 创建一个构建配置，指定TensorRT应该如何优化模型，tensorRT生成的模型只能在特定配置下运行
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    // 创建网络定义，其中createNetworkV2(1)表示采用显性batch size，新版tensorRT(>=7.0)时，不建议采用0非显性batch size
    // 因此贯穿以后，请都采用createNetworkV2(1)而非createNetworkV2(0)或者createNetwork
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

    // 构建一个模型
    /*
        Network definition:

        image
          |
        linear (fully connected)  input = 3, output = 2, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]
          |
        sigmoid
          |
        prob
    */

    // ----------------------------- 2. 输入，模型结构和输出的基本信息 -----------------------------
    const int num_input = 3;   // in_channel
    const int num_output = 2;  // out_channel
    float layer1_weight_values[] = { 1.0, 2.0, 0.5, 0.1, 0.2, 0.5 }; // 前3个给w1的rgb，后3个给w2的rgb 
    float layer1_bias_values[] = { 0.3, 0.8 };

    //输入指定数据的名称、数据类型和完整维度，将输入层添加到网络
    nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));
    nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);
    nvinfer1::Weights layer1_bias = make_weights(layer1_bias_values, 2);
    //添加全连接层
    auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias);      // 注意对input进行了解引用
    //添加激活层 
    auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); // 注意更严谨的写法是*(layer1->getOutput(0)) 即对getOutput返回的指针进行解引用

    // 将我们需要的prob标记为输出
    network->markOutput(*prob->getOutput(0));

    printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f); // 256Mib
    config->setMaxWorkspaceSize(1 << 28);
    builder->setMaxBatchSize(1); // 推理时 batchSize = 1 

    // ----------------------------- 3. 生成engine模型文件 -----------------------------
    //TensorRT 7.1.0版本已弃用buildCudaEngine方法，统一使用buildEngineWithConfig方法
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if (engine == nullptr) {
        printf("Build engine failed.\n");
        return -1;
    }

    // ----------------------------- 4. 序列化模型文件并存储 -----------------------------
    // 将模型序列化，并储存为文件
    nvinfer1::IHostMemory* model_data = engine->serialize();
    FILE* f = fopen("engine.trtmodel", "wb");
    fwrite(model_data->data(), 1, model_data->size(), f);
    fclose(f);

    // 卸载顺序按照构建顺序倒序
    model_data->destroy();
    engine->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
    printf("Done.\n");
    return 0;
}

导出ONNX

1.对于任何用到shape、size返回值的参数时，例如: tensor.view(tensor.size(0),-1)这类操作，避免直接使用tensor.size的返回值，而是加上int转换，tensor.view(int(tensor.size(O)), -1)
2.对于nn.Upsample或nn.functional.interpolate函数，使用scale_factor指定倍率，而不是使用size参数指定大小
3．对于reshape、view操作时，-1的指定请放到batch维度。其他维度可以计算出来即可。batch维度禁止指定为大于-1的明确数字
4. torch.onnx.export指定dynamic_axes参数，并且只指定batch维度，不指定其他维度。