TensorRT笔记六（构建模型并编译）

nanjono

已于 2023-04-28 13:40:48 修改

阅读量680

点赞数

文章标签：笔记人工智能 c++

于 2023-04-28 13:31:28 首次发布

本文链接：https://blog.youkuaiyun.com/nanjono/article/details/130422151

版权

文章介绍了如何使用TensorRT进行模型优化，包括定义TRTLogger类来处理日志，使用builder和config创建及配置网络，然后通过网络定义添加输入、全连接层和激活层。最后，文章阐述了如何生成engine，序列化模型并保存为文件。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、TensorRT打印日志信息

定义名为TRTLogger的类，它继承自nvinfer1::ILogger，nvinfer1是一个命名空间，包含了TensorRT库中的一系列类和函数
Severity是一个枚举类型，表示日志的严重级别（如错误、警告等）
nvinfer1::AsciiChar是一个字符类型，const* msg表示一个指向常量字符的指针，用于存储日志消息
noexcept表示这个函数不会抛出异常。override关键字表示这个函数是继承自基类的虚函数，并在派生类中进行了重写
判断严重级别，打印相应的详细信息

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);
        }
    }
};

二、TensorRT工作流程

定义网络
优化builder参数
通过builder生成engine，用于模型保存、推理等
engine通过序列化和逆序列化转化模型数据类型（转化为二进制byte文件，加快传输速率），再进一步推动模型由输入张量到输出张量的推理

在这里插入图片描述

三、代码实现结构

定义builder，config和network，其中builder为所创建的构建起，config表示创建的构建配置（指定TensorRT如何优化模型），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);

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};

    //输入指定数据的名称、数据类型和完整维度，将输入层添加到网络
    //addInput几次就有几个输入
    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标记为输出
        //模型输出节点，markOutput几次就有几个输出
	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

生成engine模型文件

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

序列化模型文件并存储

    // 将模型序列化，并储存为文件
    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");