在 Python 中，可以通过 `pynvml` 库来获取 GPU 的显存信息，并根据显存使用情况选择占用较小的显卡设备号。以下是一个完整的实现方案，包括判断显卡是否存在、获取显存使用情况，并最终选择显存占用最小的 GPU。 ### 获取所有可用 GPU 并选择显存占用最小的设备 ```python import pynvml def get_available_gpus(): """ 获取所有存在的 GPU 设备 ID """ pynvml.nvmlInit() device_count = pynvml.nvmlDeviceGetCount() available_gpus = [] for i in range(device_count): try: handle = pynvml.nvmlDeviceGetHandleByIndex(i) available_gpus.append(i) except pynvml.NVMLError: continue return available_gpus def get_gpu_memory_usage(gpu_id): """ 获取指定 GPU 的显存使用情况（单位：MB） """ pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id) mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) used = round(mem_info.used / 1024**2, 2) return used def select_least_used_gpu(): """ 获取显存占用最小的 GPU 设备号 """ available_gpus = get_available_gpus() if not available_gpus: print("没有可用的 GPU 设备。") return None gpu_usage = {} for gpu_id in available_gpus: used_mem = get_gpu_memory_usage(gpu_id) gpu_usage[gpu_id] = used_mem # 找出显存使用最少的 GPU least_used_gpu = min(gpu_usage, key=gpu_usage.get) return least_used_gpu if __name__ == "__main__": selected_gpu = select_least_used_gpu() if selected_gpu is not None: print(f"选择的 GPU 设备号：{selected_gpu}") ``` ### 功能说明 - **`get_available_gpus`**：遍历所有 GPU 索引，检查是否存在对应的 GPU 设备，返回存在的设备 ID 列表。 - **`get_gpu_memory_usage`**：获取指定 GPU 的已使用显存大小（单位为 MB）。 - **`select_least_used_gpu`**：结合前两个函数，找出所有可用 GPU 中显存占用最小的设备号。 ### 注意事项 - 需要安装 `pynvml` 库，可以通过 `pip install pynvml` 安装。 - 该代码依赖 NVIDIA 的 NVML 库，因此仅适用于安装了 NVIDIA 显卡驱动的环境。 - 若多个 GPU 的显存使用相同，`min` 函数会返回第一个找到的最小值对应的 GPU。 ### 参考文献 - 获取 GPU 显存信息的方法参考了 `pynvml` 的使用方式 [^3]。 - 显存单位转换为 MB 的计算方式参考了相关实现 [^4]。 ---