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

前言

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

01配置环境变量支持

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

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

CUDA_VISIBLE_DEVICES

02代码指定GPU设备执行

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

cudaGetDeviceCount

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

//遍历设备编号信息
intdevice;
for(device=0;device
根据查询的设备数目，GPU编号从0开始，默认情况下当前使用的设备就是编号为0的GPU设备，通过函数cudaSetDevice()可以修改运行时使用GPU设备，在初始化TensorRT之前，先通过cudaSetDevice()函数修改默认的当前设备，然后再初始化就可以把TensorRT的模型绑定到指定编号的GPU设备上推理。以我的笔记本上为例，设置当前的GPU设备，然后初始化TensorRT代码如下：
// 设置当前设备为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<getNbBindings()<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设备上执行多个模型推理的初始化代码如下：
//初始化时间标记
cudaEvent_tstart,stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);
cudaEventRecord(start,0);

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

//遍历设备编号信息
intdevice;
for(device=0;device


	

	关于延时加载

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

	

	审核编辑：汤梓红