Python运行XGBoost时如何手动指定使用的GPU设备

你可以根据自己的使用场景选以下任意一种方式配置，配置优先级为「代码内传gpu_id参数」>「代码内设置CUDA环境变量」>「启动脚本时设置CUDA环境变量」。

1. 直接通过XGBoost内置参数指定（最推荐，无副作用） #

XGBoost的GPU相关参数内置了gpu_id字段，直接填入你要使用的物理GPU编号即可，编号和nvidia-smi命令输出的GPU序号完全对应。

• 原生XGBoost训练接口示例：

import xgboost as xgb

params = {
    "tree_method": "gpu_hist", # 必须开启GPU模式的树方法参数
    "gpu_id": 1, # 此处替换为你要使用的GPU编号，比如要选2号卡就填2
    # 其余训练参数按你的需求正常填写
}
# 正常构造DMatrix、启动训练即可
model = xgb.train(params, dtrain)

• Scikit-learn风格接口（XGBClassifier/XGBRegressor等）示例：

from xgboost import XGBRegressor
# 初始化模型时直接传入gpu_id参数
reg = XGBRegressor(
    tree_method="gpu_hist",
    gpu_id=2 # 指定使用物理2号GPU
)
reg.fit(X_train, y_train)

2. 启动脚本时通过环境变量临时指定（无需改代码） #

如果你不想修改已有代码，可以在启动Python脚本时，通过CUDA_VISIBLE_DEVICES环境变量限制当前进程可见的GPU范围。设置后XGBoost会默认在可见的GPU列表里选第一个作为运行设备。
比如你要使用物理编号为3的GPU跑train.py，执行以下命令即可：

CUDA_VISIBLE_DEVICES=3 python train.py

注意：该配置仅对当前启动的进程生效，不会修改服务器全局配置，也不会影响其他用户的任务。如果需要指定多张可用GPU，用逗号分隔编号即可，例如CUDA_VISIBLE_DEVICES=1,3，此时程序内识别到的0号卡对应物理1号卡，1号卡对应物理3号卡。

3. 代码内全局设置可见GPU #

如果你希望整个Python脚本全程只能访问指定GPU，可以在导入XGBoost之前提前设置环境变量：

import os
# 限制当前进程仅能访问物理1号GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
# 之后再导入xgboost
import xgboost as xgb

# 后续正常编写训练逻辑即可，此时默认的gpu_id=0就对应物理1号卡

小提示：选卡前先执行nvidia-smi查看各张GPU的显存占用、GPU利用率，优先选择空闲资源充足的卡使用，避免影响同服务器其他用户的任务。

内容的提问来源于stack exchange，提问作者LisaBinder