Tensorrt的使用

        TensorRT(简写为TRT)是英伟达推出的针对英伟达显卡上深度学习模型部署加速的工具,使用时通常包含两个主要步骤,一是加载模型,二是执行推理,数据的前后处理有的也可以考虑放在trt模块中以提供显卡的加速能力,后续文章将会逐步展开。

        TensorRT依赖cuda和cudnn,在配置运行环境时不建议统一使用最新的cuda环境,而是根据硬件的发售时期选择发售当时推出的cuda版本或者其下一个版本

完整代码见:learn: 储存相关demo

使用TRT基本代码需要包含头文件

#include <cuda_runtime.h>

#include <NvInfer.h>

 使用TensorRT首先需要创建自己的日志类,继承并重写TRT的log函数,最基本的打印如下,也可自行添加日志级别过滤、文档保存等:

class JLogger : public nvinfer1::ILogger
{
public:
    virtual void log(Severity severity, const char* msg) noexcept override{
        printf("LOG(%d):%s\n", severity, msg);
    }
};

首先创建网络,以下代码为手动构建一个3*3的卷积网络,并保存,为方便理解每行都加了注释,最常用的做法还是直接调用onnx模型进行反序列化。

void buildModel()
{
    JLogger logger;//实例化logger,用作日志保存
    auto builder = nvinfer1::createInferBuilder(logger);
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U);//创建网络
    auto config  = builder->createBuilderConfig(); //创建编译配置
    builder->setMaxBatchSize(1);//配置batchsize
    config->setMaxWorkspaceSize(1 << 30); //分配工作空间
    //这个工作空间为cuda自己管理的内存复用空间,优化时可调小
    float kernel_weight[] = { //准备卷积核
        1, 0, 0,
        0, 1, 0,
        0, 0, 1
    };
    nvinfer1::Weights conv1_weight; //初始化卷积核权重
    nvinfer1::Weights conv1_no_bias;//初始化卷积核偏置
    conv1_weight.count = sizeof(kernel_weight) / sizeof(kernel_weight[0]);//设定权重数量
    conv1_weight.values= kernel_weight;//设定权重数据
    conv1_weight.type  = nvinfer1::DataType::kFLOAT;//设定数据类型    
    conv1_no_bias.count = 0;//设定偏置项数量
    conv1_no_bias.values= nullptr;//设定偏置数据
    conv1_no_bias.type = nvinfer1::DataType::kFLOAT;//设定偏置数据类型
    auto input =network->addInput(//网络添加输入节点
        "Image",//设定输入节点的名称
        nvinfer1::DataType::kFLOAT,//设定输入节点的数据类型
        nvinfer1::Dims4(1, 1, 3, 3));//设定输入节点的shape
    auto conv1 = network->addConvolution(
        *input,//输入节点的数据,需使用引用
        1,//输出的通道数
        nvinfer1::DimsHW(3, 3),//卷积核大小
        conv1_weight,//指定卷积权重
        conv1_no_bias//指定偏置项
    );
    //添加激活层
    auto relu1 = network->addActivation(
        *conv1->getOutput(0),   //使用卷积层的输出
        nvinfer1::ActivationType::kRELU// 使用ReLU
    );
    auto output = relu1->getOutput(0);//获取ReLU的输出
    output->setName("Predict");//设定输出节点名称为Predict
    network->markOutput(*output);//将output标记为网络输出

    auto engine = builder->buildEngineWithConfig(*network, *config);//构建引擎
    auto host_memory = engine->serialize();//获取引擎指针
    const char model_save_path[] = "cnn.engine";
    save_to_file(model_save_path, host_memory->data(), host_memory->size());
    host_memory->destroy();
    engine->destroy();
    config->destroy();
    network->destroy();
    builder->destroy();
    cout << "save to " << model_save_path << endl;
}

接下来可以调用上一步构建的网络,执行推理,

void model_inference()
{
    JLogger logger;//建立日志类
    cudaStream_t stream = nullptr;//创建流
    cudaStreamCreate(&stream);
    cudaSetDevice(0);//选择显卡
    const char engine_path[] = "/data/learn/learn/Tensorrt/workspace/cnn.engine";
    auto model_data = load_from_file(engine_path);//加载引擎文件
    if (model_data.empty()){
        printf("model %s load failed.", engine_path);
    }
    auto runtime = nvinfer1::createInferRuntime(logger);//创建运行时对象
    auto engine = runtime->deserializeCudaEngine(model_data.data(), model_data.size());//反序列化引擎
    auto context = engine->createExecutionContext();//创建上下文
    //获取输入输出节点相关信息
    // int nbindings = engine->getNbBindings();//获取输入输出节点总数
    // for (int i=0; i<nbindings; i++)
    // {
    //     auto dims = engine->getBindingDimensions(i);//获取节点shape
    //     auto name = engine->getBindingName(i);//获取节点名称
    //     auto type = engine->bindingIsInput(i) ? "input" : "output";//判断是否为输入节点
    //     cout << "binging" << name << i << "dim is" << format_dim(dims) << " type is " << type << endl;
    // }
    //准备输入输出数据及内存
    float* output_data_device = nullptr;//准备输出数据内存
    cudaMalloc(&output_data_device, sizeof(float));
    float output_data[1];
    float input_data[] = {//准备输入数据
        1,2,3,
        4,5,6,
        7,8,9
    };
    size_t input_size = sizeof(input_data);//获取输入数据大小
    float* input_data_device = nullptr;
    cudaMalloc(&input_data_device, input_size);//分配输入数据的显存
    cudaMemcpyAsync(input_data_device, input_data, input_size, cudaMemcpyHostToDevice, stream);//输入数据异步拷贝
    void* bingdings_device_pointer[] = {input_data_device, output_data_device};//将输入输出的地址按顺序储存
    //执行推理
    bool finished = context->enqueueV2( bingdings_device_pointer, stream, nullptr);
    if(!finished)
    {
        cout << "Enqueue failed." << endl;
    }
    cudaMemcpyAsync(output_data, output_data_device, sizeof(float), cudaMemcpyDeviceToHost, stream);//将执行结果传回cpu
    cudaStreamSynchronize(stream);//等待数据传输完毕
    cout << "output data size" << sizeof(output_data) << endl;
    cout << "Enqueue succeed, result is " << output_data[0] << endl;
    //释放内存
    cudaFree(input_data_device);
    cudaFree(output_data_device);
    context->destroy();
    engine->destroy();
    runtime->destroy();
    cudaStreamDestroy(stream);
}

动态输入大小:

auto profile = builder.createOptimizationProfile();

profile->setDimensions("foo", OptProfileSelector::kMIN, Dims4(n, c, h, w)));

config.addOpeimizationProfile(profile);

其中OptProfileSelector::kMIN表示最小输入尺寸,kOPT为最常用输入尺寸,kMAX为最大输入尺寸,可根据实际使用情况做调整,使用时只需要在范围之间就能运行。

  • 9
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值