TensorRT | 在多个GPU中指定推理设备

2023-08-22 13:41:17 浏览数 (1)
前言

说实话，之前我在笔记本上都一直都是只有一块N卡，所以没有过多关注过这个问题。然而昨天有个人问我，TensorRT怎么在多个GPU中指定模型推理GPU设备？我查了一下，发现官方有几个不同的解决方案，个人总结了一下，主要的做法有两种。

01

配置环境变量支持

该方法的好处是不需要修改代码，通过配置环境变量就可以实现指定的GPU运行，缺点是缺乏灵活性，特别是想切换不同GPU实现模型推理的时候，这个方法就弊端就比较明显。

CUDA编程中支持的指定GPU设备的环境变量为：

代码语言：javascript复制
CUDA_VISIBLE_DEVICES

通过该系统的环境变量可以设置指定的单个GPU编号或者多个GPU编号合集，然后在程序测试与调试环境中使用。通过这种方式指定GPU编号执行模型推理，就无需修改代码，实现在单一指定的GPU上运行TensorRT推理程序。

02

代码指定GPU设备执行

一台机器上可能有多个GPU设备，通过CUDA编程可以查询机器上所有的GPU设备，查询这些设备的属性以及决定使用哪个GPU设备作为当前设备。

代码语言：javascript复制
cudaGetDeviceCount

该函数可以查询到当前机器上GPU设备数目，然后遍历查询每个GPU设备的属性。官方教程给出的代码如下：

代码语言：javascript复制
// 查询设备数目
int deviceCount;
cudaGetDeviceCount(&deviceCount);

// 遍历设备编号信息
int device;
for (device = 0; device < deviceCount;   device) {
    cudaDeviceProp deviceProp;
    cudaGetDeviceProperties(&deviceProp, device);
    printf("Device %d has compute capability %d.%d.n",
        device, deviceProp.major, deviceProp.minor);
}

根据查询的设备数目，GPU编号从0开始，默认情况下当前使用的设备就是编号为0的GPU设备，通过函数cudaSetDevice()可以修改运行时使用GPU设备，在初始化TensorRT之前，先通过cudaSetDevice()函数修改默认的当前设备，然后再初始化就可以把TensorRT的模型绑定到指定编号的GPU设备上推理。以我的笔记本上为例，设置当前的GPU设备，然后初始化TensorRT代码如下：

代码语言：javascript复制
// 设置当前设备为GPU 0
cudaSetDevice(0);// 初始化TensorRT
this->runtime = createInferRuntime(gLogger);
assert(this->runtime != nullptr);
this->engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(this->engine != nullptr);
this->context = engine->createExecutionContext();
assert(this->context != nullptr);
delete[] trtModelStream;

// do more thing here// insert query input and output layers information
// 创建GPU显存输入/输出缓冲区
std::cout << " input/outpu : " << engine->getNbBindings() << std::endl;
cudaMalloc(&buffers[input_index], this->input_h * this->input_w * 3 * sizeof(float));
cudaMalloc(&buffers[2], this->output_h *this->output_w * sizeof(float));
cudaMalloc(&buffers[1], 32 *25600 * sizeof(float));

// 创建临时缓存输出
prob.resize(output_h * output_w);
mprob.resize(32 * 25600);

// 创建cuda流
cudaStreamCreate(&stream);

在多个GPU设备上执行多个模型推理的初始化代码如下：

代码语言：javascript复制
// 初始化时间标记
cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);
cudaEventRecord(start, 0);

// 查询设备数目
int deviceCount;
cudaGetDeviceCount(&deviceCount);

// 遍历设备编号信息
int device;
for (device = 0; device < deviceCount;   device) {
    cudaDeviceProp deviceProp;
    cudaGetDeviceProperties(&deviceProp, device);
    printf("Device %d has compute capability %d.%d.n",
        device, deviceProp.major, deviceProp.minor);
}

// Set GPU 0 as current
cudaSetDevice(0);            
cudaStream_t s0;
cudaStreamCreate(&s0);
void* p0[1];
size_t size = 1024 * sizeof(float);
cudaMalloc(p0[0], size);
// initialization TensorRT here on GPU 0 

// Set GPU 1 as current
cudaSetDevice(1);            
cudaStream_t s1;
cudaStreamCreate(&s1);
// initialization TensorRT here on GPU 1

// 计算执行时间
cudaEventRecord(stop, 0);
cudaEventSynchronize(stop);
float elapsedTime;
cudaEventElapsedTime(&elapsedTime, start, stop);
printf("time to consume: %3.1f ms n", elapsedTime);

// 销毁
cudaEventDestroy(start);
cudaEventDestroy(stop);

关于延时加载

TensorRT8.6支持CUDA Lazy Loading(延时加载)，开发者文档上说这种方式可以有效降低GPU显存与内存使用，加速初始化，节省模型初始化时间，可以通过环境变量配置实现延时加载支持，相关环境变量为：

代码语言：javascript复制
CUDA_MODULE_LOADING=LAZY

参考资料：

代码语言：javascript复制
https://developer.nvidia.com/blog/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/
https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#multi-device-system

掌握TensorRT8.6 C 如何部署图像分类，对象检测，实例分割，语义分割主流模型，自定义脚本一键INT8量化模型，使用OpenCV CUDA加速图像预处理等各种工程化部署推理技巧，实现推理部署的工程化封装支持

gpu 部署 函数 模型 配置

0 人点赞