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

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

01 配置环境变量支持

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

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

复制代码
CUDA_VISIBLE_DEVICES

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

02 代码指定GPU设备执行

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

复制代码
cudaGetDeviceCount

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

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

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

关于延时加载

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

复制代码
CUDA_MODULE_LOADING=LAZY

参考资料:

复制代码
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
相关推荐
驭白.几秒前
从订单到行驶:构建新能源汽车产品全生命周期数据链
人工智能·汽车·制造·数字化·制造业·新能源汽车
一瞬祈望5 分钟前
⭐ 深度学习入门体系(第 15 篇): 从 RNN 到 LSTM:为什么深度网络需要“记忆能力”?
rnn·深度学习·lstm
Watermelo6175 分钟前
探究TOON的价值边界:比JSON更优的大模型友好数据格式?
数据结构·人工智能·语言模型·自然语言处理·数据挖掘·数据分析·json
LeeeX!6 分钟前
基于YOLO11实现明厨亮灶系统实时检测【多场景数据+模型训练、推理、导出】
深度学习·算法·目标检测·数据集·明厨亮灶
模型启动机6 分钟前
DeepSeek-OCR是「长文本理解」未来方向?中科院新基准VTCBench给出答案
人工智能·ai·大模型·ocr
德育处主任Pro8 分钟前
『n8n』环境搭建
人工智能·aigc·n8n
GIS学姐嘉欣8 分钟前
新中地GIS开发校企联合实训 | 山东理工大学优秀作品(1)广州市智慧城市监测平台
人工智能·智慧城市
我是宝库10 分钟前
Turnitin系统查英文论文AI率,准吗?
人工智能·论文ai率·turnitin系统·英文查重·aigc检测·英文ai率·英文查ai率
知乎的哥廷根数学学派10 分钟前
基于高阶统计量引导的小波自适应块阈值地震信号降噪算法(MATLAB)
网络·人工智能·pytorch·深度学习·算法·机器学习·matlab
墨北小七11 分钟前
CNN深度学习模型在小说创作领域的应用
人工智能·深度学习·cnn