前言

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

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

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

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

CUDA_VISIBLE_DEVICES

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

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

cudaGetDeviceCount

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

//遍歷設(shè)備編號(hào)信息
intdevice;
for(device=0;device
根據(jù)查詢(xún)的設(shè)備數(shù)目，GPU編號(hào)從0開(kāi)始，默認(rèn)情況下當(dāng)前使用的設(shè)備就是編號(hào)為0的GPU設(shè)備，通過(guò)函數(shù)cudaSetDevice()可以修改運(yùn)行時(shí)使用GPU設(shè)備，在初始化TensorRT之前，先通過(guò)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);

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

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


	

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

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

	

	審核編輯：湯梓紅