如何檢查 XGBoost 是否使用 GPU

Question

我正在編寫一個 pytest 文件來檢查我的機器學習庫是否使用 GPU。 對於 Tensorflow 我可以用tf.config.list_physical_devices()檢查這個。 對於 XGBoost，到目前為止，我在運行我的軟件時通過查看 GPU 利用率 ( nvdidia-smi ) 對其進行了檢查。 但是我怎樣才能在一個簡單的測試中檢查呢？ 類似於我對 Tensorflow 進行的測試可以做到。

import pytest
import tensorflow as tf
import xgboost

# Marking all tests to be GPU dependent
pytestmark = pytest.mark.gpu

def test_tf_finds_gpu():
    """Check if Tensorflow finds the GPU."""
    assert tf.config.list_physical_devices("GPU")

def test_xgb_finds_gpu():
    """Check if XGBoost finds the GPU."""
    ...
    # What can I write here?

Answer 1

我使用的測試方法是使用tree_method="gpu_hist"運行的。 根據我無法確定的情況，如果找不到 GPU，則會引發錯誤或打印警告。

因此，如果找不到 GPU，以下測試將通過以下兩種方式之一捕獲它：

• 在xgb_model.fit(X, y)上引發XGBoostError
• 在xgb_model.fit(X, y)上打印警告。 這將由 pytest 提供的capsys夾具captured.out ，並且 capture.out 或captured.err不會為空。 因此，其中一個斷言將引發AssertionError 。

from sklearn.datasets import load_boston

def test_xgb_finds_gpu(capsys):
    """Check if XGBoost finds the GPU."""
    boston = load_boston()
    X = boston["data"]
    y = boston["target"]
    xgb_model = xgb.XGBRegressor(
        # If there is no GPU, the tree_method kwarg will cause either
        # - an error in `xgb_model.fit(X, y)` (seen with pytest) or
        # - a warning printed to the console (seen in Spyder)
        # It's unclear which of the two happens under what circumstances.
        tree_method="gpu_hist"
    )
    xgb_model.fit(X, y)
    # Check that no warning was printed.
    captured = capsys.readouterr()
    assert captured.out == ""
    assert captured.err == ""

我認為可以通過對X和y使用更小的 arrays 來加快這個測試，但是實現這個需要我太多的時間，因為沒有 GPU 的測試只需要幾秒鍾，而使用 Z52F9EC213CA077D32Z 則不到一秒鍾。