python写个检测本机GPU的程序

要使用 Python 编写检测本机 GPU 的程序，可根据不同的 GPU 厂商和框架来实现。以下分别给出使用 NVIDIA GPU 的 CUDA 以及通用的 PyTorch 和 TensorFlow 框架的示例代码。 ### 使用 NVIDIA CUDA 若使用 NVIDIA GPU，可借助 `pynvml` 库来检测 GPU 信息。需先安装该库： ```bash pip install pynvml ``` 示例代码如下： ```python import pynvml # 初始化 NVML pynvml.nvmlInit() # 获取 GPU 设备数量 device_count = pynvml.nvmlDeviceGetCount() print(f"本机检测到的 GPU 数量: {device_count}") # 遍历每个 GPU 设备并打印信息 for i in range(device_count): handle = pynvml.nvmlDeviceGetHandleByIndex(i) info = pynvml.nvmlDeviceGetMemoryInfo(handle) name = pynvml.nvmlDeviceGetName(handle) print(f"GPU {i} 信息:") print(f" 名称: {name.decode('utf-8')}") print(f" 总显存: {info.total / 1024 ** 3:.2f} GB") print(f" 已使用显存: {info.used / 1024 ** 3:.2f} GB") print(f" 空闲显存: {info.free / 1024 ** 3:.2f} GB") # 关闭 NVML pynvml.nvmlShutdown() ``` ### 使用 PyTorch 若使用 PyTorch 框架，可使用以下代码检测 GPU 是否可用以及获取 GPU 信息： ```python import torch # 检查 GPU 是否可用 if torch.cuda.is_available(): print("GPU 可用") # 获取 GPU 设备数量 device_count = torch.cuda.device_count() print(f"本机检测到的 GPU 数量: {device_count}") # 打印每个 GPU 的名称 for i in range(device_count): print(f"GPU {i} 名称: {torch.cuda.get_device_name(i)}") else: print("GPU 不可用") ``` ### 使用 TensorFlow 若使用 TensorFlow 框架，可使用以下代码检测 GPU 是否可用： ```python import tensorflow as tf # 检查 GPU 是否可用 gpus = tf.config.list_physical_devices('GPU') if gpus: print("GPU 可用") print(f"本机检测到的 GPU 数量: {len(gpus)}") for gpu in gpus: print(f"GPU 信息: {gpu}") else: print("GPU 不可用") ```

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考