【TRT】内存管理封装

1. host 和 device 内存封装

问题提出：

float* input_data_host = nullptr;
float* input_data_device = nullptr;
checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));
checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));

以上可以看到 host 和 devie memory 经常成对出现，所以这里我们可以将二者进行封装。

• 由 MixMemory 来负责申请Host 和 Device 内存
• 复用MixMemory中的GPU和CPU内存，避免反复申请
• 当申请的内存大于目前的内存时，释放当前内存，重新申请内存

 MixMemory input_data;
 float* input_data_host   = input_data.cpu<float>(input_numel);
 float* input_data_device = input_data.gpu<float>(input_numel);

2. Tensor封装

• 内存管理，可以使用上面的MixMemory解决
• 内存的复用，可以使用MixMemory解决
• 内存的拷贝，比如 CPU --> GPU, GPU --> CPU
• 索引的计算，比如有一个5D的Tensor(B,D,C,H,W) 此时需要获取一个特定位置的元素。

2.1 内存拷贝

• 定义内存的标记，表示内存当前最新的内容在哪里(GPU |CPU | Init)
• 懒分配原则，当需要使用时，才会考虑分配内存
• 获取内存地址，即表示：想拿到最新的数据。

  •   比如tensor.cpu 表示想拿到最新的数据，并且把它放在cpu 上。
    

  •   比如tensor.gpu 表示想拿到最新的数据，并且把它放在gpu上。
    

不封装的代码

int input_batch = 1;
int input_channel = 3;
int input_height = 224;
int input_width = 224;
int input_numel = input_batch * input_channel * input_height * input_width;
float* input_data_host = nullptr;
float* input_data_device = nullptr;
checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));
checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));
// 对应于pytorch的代码部分
    cv::resize(image, image, cv::Size(input_width, input_height));
    int image_area = image.cols * image.rows;
    unsigned char* pimage = image.data;
    float* phost_b = input_data_host + image_area * 0;
    float* phost_g = input_data_host + image_area * 1;
    float* phost_r = input_data_host + image_area * 2;
    for(int i = 0; i < image_area; ++i, pimage += 3){
        // 注意这里的顺序rgb调换了
        *phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];
        *phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];
        *phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];
    }
    ///
    checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel * sizeof(float), cudaMemcpyHostToDevice, stream));

    // 3x3输入，对应3x3输出
    const int num_classes = 1000;
    float output_data_host[num_classes];
    float* output_data_device = nullptr;
    checkRuntime(cudaMalloc(&output_data_device, sizeof(output_data_host)));

    // 明确当前推理时，使用的数据输入大小
    auto input_dims = execution_context->getBindingDimensions(0);
    input_dims.d[0] = input_batch;

    // 设置当前推理时，input大小
    execution_context->setBindingDimensions(0, input_dims);
    float* bindings[] = {input_data_device, output_data_device};
    bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);
    checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream));
    checkRuntime(cudaStreamSynchronize(stream));

封装后的代码

cudaStream_t stream = nullptr;
checkRuntime(cudaStreamCreate(&stream));
// 申请一个 [batch, channel, height, width] 类型为 Float的tensor
TRT::Tensor input_data({input_batch, input_channel, input_height, input_width}, TRT::DataType::Float);
// 为input关联stream，使得在同一个pipeline中执行复制操作
input_data.set_stream(stream);
///
    // image to float
    auto image = cv::imread("dog.jpg");
    float mean[] = {0.406, 0.456, 0.485};
    float std[]  = {0.225, 0.224, 0.229};

    // 对应于pytorch的代码部分
    cv::resize(image, image, cv::Size(input_width, input_height));
    image.convertTo(image, CV_32F);

    // 利用opencv mat的内存地址引用，实现input与mat的关联，然后利用split函数一次性完成mat到input的复制
    cv::Mat channel_based[3];
    for(int i = 0; i < 3; ++i)
        // 注意这里 2 - i是实现bgr -> rgb的方式
        // 这里cpu提供的参数0是表示batch的索引是0，第二个参数表示通道的索引，因此获取的是0, 2-i通道的地址
        // 而tensor最大的好处就是帮忙计算索引，否则手动计算就得写很多代码
        // 这里 取tensort的 第 0 个batch， channel为 2-i 的数据构建Mat
        // 这里使用了OpenCV的 "Header" 模式，即Mat 的数据和这里tensor的数据指向同一块内存 
        channel_based[i] = cv::Mat(input_height, input_width, CV_32F, input_data.cpu<float>(0, 2-i));
    // 将rgb 的数据分为 rrr ggg bbb 的tensor
    cv::split(image, channel_based);
// 利用opencv的mat操作加速减去均值和除以标准差
  for(int i = 0; i < 3; ++i)
       channel_based[i] = (channel_based[i] / 255.0f - mean[i]) / std[i];

// 如果不写，input_data.gpu获取gpu地址时会自动进行复制
// 目的就是把内存复制变为隐式进行
input_data.to_gpu();

// 3x3输入，对应3x3输出
const int num_classes = 1000;
TRT::Tensor output_data({input_batch, num_classes}, TRT::DataType::Float);
output_data.set_stream(stream);

// 明确当前推理时，使用的数据输入大小
auto input_dims = execution_context->getBindingDimensions(0);
input_dims.d[0] = input_batch;

execution_context->setBindingDimensions(0, input_dims);
float* bindings[] = {input_data.gpu<float>(), output_data.gpu<float>()};
bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);
checkRuntime(cudaStreamSynchronize(stream));

// 当获取cpu地址的时候，如果数据最新的在gpu上，就进行数据复制，然后再返回cpu地址
float* prob = output_data.cpu<float>();