TensorRT学习系列一:使用自定义网络构建和运行TensorRT模型

2023-12-16 22:06:53


一、前言

随着深度学习技术的发展,许多实际应用都采用了神经网络模型来处理复杂任务,如图像识别、语音识别、自然语言处理等。然而,在部署这些模型时,如何实现在生产环境中高效且准确地执行推理是一个重要问题。这就需要利用专门针对推理优化的工具和技术,其中NVIDIA的TensorRT就是这样一款高性能的深度学习推理优化器。

TensorRT是一个由NVIDIA开发的高性能推理平台,旨在加速深度学习模型的推断阶段,并最大限度地提高GPU的利用率。它通过一系列优化技术,包括模型剪枝、量化以及算法选择等,能够在保持模型精度的同时,显著提升模型在生产环境中的运行效率。
TensorRT广泛应用于各种场景,例如自动驾驶、医疗影像分析、云计算服务、实时推荐系统等,特别是在需要实时响应或处理大量数据流的情况下,其优势尤为明显。

本文是“TensorRT学习系列”之一,我们将专注于如何使用自定义网络构建和运行TensorRT模型。我们将详细介绍如何将训练好的模型导入TensorRT,并对其进行优化以实现高效的推理。此外,本文还将特别关注TensorRT构建和运行时的新旧API对比说明,帮助读者更好地理解TensorRT的工作原理,以便在实际项目中做出最佳实践的选择。

在这个过程中,我们不仅会探讨TensorRT的基本使用方法,还会涉及到性能优化策略和常见问题的解决方案。希望通过本文的学习,你能掌握如何有效地利用TensorRT提升深度学习模型的推理性能。

关于运行环境的搭建和TensorRT的安装,请参考我另一篇博客:Ubuntu下安装ONNX、ONNX-TensorRT、Protobuf和TensorRT,本文不再详述。


二、使用TensorRT构建(build)引擎

我们使用的TensorRT 8.6-GA版,比这之前的统称旧版,部分代码可能在这之前的某个版本声明将要废弃(deprecated),并不一定是这一版开始声明将要废弃的。这里不单独声明,具体可以查询TensorRT的官方api文档。另外,这些声明将要废弃的代码在这一版还是可以运行的。
下面是TensorRT build engine的主要过程

  1. 创建logger
  2. 创建builder
  3. 创建network (builder->network)
  4. 配置参数 (builder->config)
  5. 生成engine和序列化保存 (builder->serialized_engine(network, config))
  6. 释放资源 (delete)

我们分别进行代码说明:

说明:我们在代码中会有#define oldversion 0,这个oldversion的宏就是指代之前的版本写法。
另外,下面很多代码里有缩进,那是因为本身是在main函数里。

1. 创建logger

这个写法相对固定,主要是创建builder的时候需要,所以写在前面

class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char *msg) noexcept override {
        if (severity != Severity::kINFO) {
            std::cout << msg << std::endl;
        }
    }
} gLogger;

2. 创建builder

这个写法相对固定,新旧版本之间没有太大差异。

	nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger);

3. 创建network

这是一种指定显性batch的写法,也可以直接指定参数为1。注意使用的是V2。

	auto explictBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);    
	nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explictBatch);

我们讲自定义网络,手动添加输入层,全连接层和激活层作为案例说明。

    const int input_size = 3;
    const int output_size = 2;
    // 标记输入层和创建输入数据
    nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});
    const float *fc1_weight_data = new float[input_size * output_size]{0.1,0.2,0.3,0.4,0.5,0.6};
    const float *fc1_bias_data = new float[2]{0.1,0.5};
    nvinfer1::Weights fc1_weight{nvinfer1::DataType::kFLOAT, fc1_weight_data, input_size * output_size}; 
    nvinfer1::Weights fc1_bias{nvinfer1::DataType::kFLOAT, fc1_bias_data, output_size};
// 下面是比较核心的部分,是新旧版本目前看来差异是大的部分
#if oldversion
    // 设定最大batch size
    builder->setMaxBatchSize(1);
    nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);
    nvinfer1::IActivationLayer *sigmoid = network->addActivation(*fc1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); 
#else
    // 将输入张量转换为2D,以进行矩阵乘法
    nvinfer1::IShuffleLayer* shuffleLayer = network->addShuffle(*input);
    shuffleLayer->setReshapeDimensions(nvinfer1::Dims2{input_size, 1});
    // 创建权重矩阵的常量层
    nvinfer1::IConstantLayer* weightLayer = network->addConstant(nvinfer1::Dims2{output_size, input_size}, fc1_weight);
    // 添加矩阵乘法层
    nvinfer1::IMatrixMultiplyLayer* matMulLayer = network->addMatrixMultiply(
        *weightLayer->getOutput(0), nvinfer1::MatrixOperation::kNONE,
        *shuffleLayer->getOutput(0), nvinfer1::MatrixOperation::kNONE
    );
    // 添加偏置
    nvinfer1::IConstantLayer* biasLayer = network->addConstant(nvinfer1::Dims2{output_size, 1}, fc1_bias);
    nvinfer1::IElementWiseLayer* addBiasLayer = network->addElementWise(
        *matMulLayer->getOutput(0), *biasLayer->getOutput(0), nvinfer1::ElementWiseOperation::kSUM);
    // 添加激活层
    nvinfer1::IActivationLayer* sigmoid = network->addActivation(*addBiasLayer->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
#endif
		// 标记输出层
    sigmoid->getOutput(0)->setName("output");
    network->markOutput(*sigmoid->getOutput(0));

创建全连接层和激活层的方法在新旧版本中差异较大,需要多关注。

4. 配置参数

createBuilderConfig接口被用来指定TensorRT应该如何优化模型。这里应该注意,在旧版builder->setMaxBatchSize(1);这个代码一定要有,否则运算结果不对。设置最大工作区大小,相当于设置内存池,新旧版本api不同。

#if oldversion
    builder->setMaxBatchSize(1);
#endif
    nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 设置最大工作区大小,这个就是api的不同
#if oldversion
    config->setMaxWorkspaceSize(1 << 28);
#else
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1<<28);
#endif

5. 构建并序列化引擎并保存

这个在旧版本中是分两步,构建和序列化是分开的,在新版中合为一个api。这所以要序列化,主要是构建过程比较长,一般序列化后要保存成文件,这里不作过多介绍。

// 构建并序列化引擎
#if oldversion
    nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
    if (!engine) {
        std::cerr << "build engine failed" << std::endl;
		// 将之前堆内存分配的变量释放
        return -1;
    }
    nvinfer1::IHostMemory *serialized_engine = engine->serialize();
#else
// 使用 buildSerializedNetwork 构建并序列化网络
    nvinfer1::IHostMemory* serialized_engine = builder->buildSerializedNetwork(*network, *config);
    if (!serialized_engine) {
        std::cerr << "build engine failed" << std::endl;
		// 将之前堆内存分配的变量释放
        return -1;
    }
#endif

6. 释放资源

这块比较简单,不详述。就是注意一下,在旧版本中使用对象的destory方法,新版本直接delete。再一个就是注意按照构造的顺序反过来进行释放。

#if oldversion
	...
	builder.destroy();
#else
	...
	delete builder;
#endif

三、使用TensorRT运行(runtime)模型

上面构建好了引擎,我们需要运行模型进行推理,下面是runtime的主要过程

  1. 创建一个runtime对象
  2. 反序列化生成engine
  3. 创建一个执行上下文
  4. 填充数据
  5. 执行推理
  6. 释放资源
    我们分别进行代码说明:

logger和上部分一样,不单独列出来了
下面我部分采取了智能指针的写法,自动释放资源。

1. 创建一个runtime对象

标准写法。

	// 创建一个runtime对象
	auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(gLogger));

2. 反序列化生成engine

loadData是自己实现从文件读取序列化数据,不再详述。

    // 反序列化生成engine
    std::string serialized_engine_file = argv[1];
    std::vector<char> engine_data = loadData(serialized_engine_file);
    auto engine = std::unique_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));
    if (engine == nullptr) {
        std::cout << "deserializeCudaEngine failed" << std::endl;
        // 因为采用的是智能指针,不必再手动释放资源。
        return -1;
    }

3. 创建一个执行上下文

    auto context = std::unique_ptr<nvinfer1::IExecutionContext>(engine->createExecutionContext());
    if (context == nullptr) {
        std::cout << "createExecutionContext failed" << std::endl;
        return -1;
    }
#endif

4. 填充数据

这块是cuda编程非常常规的写法,这里不再详细说明。

    // 输入数据
    // float host_input_data = new float[input_size]{1,2,3};
    std::unique_ptr<float[]> host_input_data(new float[input_size]{2,4,8});
    int host_intput_size = input_size * sizeof(float);
    float *device_input_data = nullptr;
    // 输出数据
    // float host_output_data = new float[output_size]{0};
    std::unique_ptr<float[]> host_output_data(new float[output_size]{0});
    int host_output_size = output_size * sizeof(float);
    float *device_output_data = nullptr;
    // 创建 CUDA 流
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    // 申请device内存
    cudaMalloc((void **)&device_input_data, host_intput_size);
    cudaMalloc((void**)&device_output_data, host_output_size);
    cudaMemcpyAsync(device_input_data, host_input_data.get(), host_intput_size, cudaMemcpyHostToDevice, stream);

5. 执行推理

我们注意到在旧版本中使用的是enqueueV2,但在新版本中用的是executeV2,而不是enqueueV3。根据网上资料enqueueV3和executeV2作用是相同的,但是只接受一个stream作为参数,但没有实践成功。我一直想研究一个enqueueV3,因为我开始感觉这是最正规的做法,后来看了TensorRT官方sample里面的案例,用的就是executeV2。所以executeV2才是目前最正规的做法。

    // 准备绑定缓冲区
    void* bindings[] = {device_input_data, device_output_data};
#if oldversion
    bool status = context->enqueueV2((void**)bindings, stream, nullptr);
#else
    bool status = context->executeV2(bindings);
#endif

6. 释放资源


四、总结

本文深入探讨了如何使用自定义网络构建和运行TensorRT模型。首先,介绍了TensorRT的基础知识,包括其基本概念和应用场景。接着,详细讲解了如何将自定义网络转换为TensorRT模型,并在此过程中进行了新旧API的对比说明。随后,展示了如何使用TensorRT运行模型并获取结果。在这些过程中, 列举了一些常见的问题及其解决方案,希望能帮助读者在实际操作中避免这些问题。

五、参考

我在github上有具体的实现,希望给读者以参考。后续的关于TensorRT的分享也会有部分在这个项目中实现。
项目github

文章来源:https://blog.csdn.net/m0_51661400/article/details/135035527
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。