TensorRT怎么在多個(gè)GPU中指定推理設(shè)備

前言

說實(shí)話，之前我在筆記本上都一直都是只有一塊N卡，所以沒有過多關(guān)注過這個(gè)問題。然而昨天有個(gè)人問我，TensorRT怎么在多個(gè)GPU中指定模型推理GPU設(shè)備？我查了一下，發(fā)現(xiàn)官方有幾個(gè)不同的解決方案，個(gè)人總結(jié)了一下，主要的做法有兩種。

01配置環(huán)境變量支持

該方法的好處是不需要修改代碼，通過配置環(huán)境變量就可以實(shí)現(xiàn)指定的GPU運(yùn)行，缺點(diǎn)是缺乏靈活性，特別是想切換不同GPU實(shí)現(xiàn)模型推理的時(shí)候，這個(gè)方法就弊端就比較明顯。

CUDA編程中支持的指定GPU設(shè)備的環(huán)境變量為：

CUDA_VISIBLE_DEVICES

通過該系統(tǒng)的環(huán)境變量可以設(shè)置指定的單個(gè)GPU編號(hào)或者多個(gè)GPU編號(hào)合集，然后在程序測試與調(diào)試環(huán)境中使用。通過這種方式指定GPU編號(hào)執(zhí)行模型推理，就無需修改代碼，實(shí)現(xiàn)在單一指定的GPU上運(yùn)行TensorRT推理程序。

02代碼指定GPU設(shè)備執(zhí)行

一臺(tái)機(jī)器上可能有多個(gè)GPU設(shè)備，通過CUDA編程可以查詢機(jī)器上所有的GPU設(shè)備，查詢這些設(shè)備的屬性以及決定使用哪個(gè)GPU設(shè)備作為當(dāng)前設(shè)備。

cudaGetDeviceCount

該函數(shù)可以查詢到當(dāng)前機(jī)器上GPU設(shè)備數(shù)目，然后遍歷查詢每個(gè)GPU設(shè)備的屬性。官方教程給出的代碼如下：

//查詢?cè)O(shè)備數(shù)目
intdeviceCount;
cudaGetDeviceCount(&deviceCount);

//遍歷設(shè)備編號(hào)信息
intdevice;
for(device=0;device
根據(jù)查詢的設(shè)備數(shù)目，GPU編號(hào)從0開始，默認(rèn)情況下當(dāng)前使用的設(shè)備就是編號(hào)為0的GPU設(shè)備，通過函數(shù)cudaSetDevice()可以修改運(yùn)行時(shí)使用GPU設(shè)備，在初始化TensorRT之前，先通過cudaSetDevice()函數(shù)修改默認(rèn)的當(dāng)前設(shè)備，然后再初始化就可以把TensorRT的模型綁定到指定編號(hào)的GPU設(shè)備上推理。以我的筆記本上為例，設(shè)置當(dāng)前的GPU設(shè)備，然后初始化TensorRT代碼如下：
// 設(shè)置當(dāng)前設(shè)備為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
//創(chuàng)建GPU顯存輸入/輸出緩沖區(qū)
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));

//創(chuàng)建臨時(shí)緩存輸出
prob.resize(output_h*output_w);
mprob.resize(32*25600);

//創(chuàng)建cuda流
cudaStreamCreate(&stream);
 在多個(gè)GPU設(shè)備上執(zhí)行多個(gè)模型推理的初始化代碼如下：
//初始化時(shí)間標(biāo)記
cudaEvent_tstart,stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);
cudaEventRecord(start,0);

//查詢?cè)O(shè)備數(shù)目
intdeviceCount;
cudaGetDeviceCount(&deviceCount);

//遍歷設(shè)備編號(hào)信息
intdevice;
for(device=0;device


	

	關(guān)于延時(shí)加載

	  TensorRT8.6支持CUDA Lazy Loading(延時(shí)加載)，開發(fā)者文檔上說這種方式可以有效降低GPU顯存與內(nèi)存使用，加速初始化，節(jié)省模型初始化時(shí)間，可以通過環(huán)境變量配置實(shí)現(xiàn)延時(shí)加載支持，相關(guān)環(huán)境變量為：
CUDA_MODULE_LOADING=LAZY

	

	審核編輯：湯梓紅