1. 项目概述为什么需要查看并指定GPU在深度学习、科学计算或者大规模数据处理的项目里GPU图形处理器早已不是游戏玩家的专属。它凭借其强大的并行计算能力成为了加速模型训练和复杂运算的“核动力引擎”。然而当你兴冲冲地准备跑一个大型模型时可能会遇到几个尴尬的场景服务器上有多块GPU但你的代码默认只用了一块其他几块在“围观”或者你和其他人共用一台服务器你的任务把别人的任务给挤掉了导致冲突。这时候学会用Python代码“侦察”可用的GPU并精确地“指挥”你的任务在指定的GPU上运行就从一个加分项变成了必备技能。这不仅仅是运行一个命令那么简单。它背后涉及到对计算资源的有效管理、对任务运行环境的精确控制以及对多GPU系统工作方式的理解。无论是个人开发者在一台多卡工作站上调试模型还是团队在共享的GPU集群上部署任务掌握这项技能都能让你事半功倍避免资源浪费和潜在冲突。接下来我将从一个实践者的角度带你一步步拆解这个需求从原理到实操再到避坑指南让你彻底搞懂。2. 核心需求解析与方案选型2.1 需求拆解我们到底要做什么看似简单的“查看并指定GPU”其实可以拆解为三个递进的核心动作探测Probe程序启动时自动识别当前系统中有哪些GPU设备以及它们的关键状态如是否空闲、显存占用、计算负载等。这相当于给你的代码装上了“雷达”。选择Select基于探测到的信息结合你的任务需求例如需要大显存或者需要避开繁忙的卡制定一个选择策略决定使用哪一块或哪几块GPU。绑定Bind将你的Python进程以及进程中用到的深度学习框架如PyTorch、TensorFlow的计算任务强制绑定到你所选定的GPU设备上确保任务不会“跑偏”。2.2 主流工具选型为什么是它们在Python生态中完成这些动作主要依赖两个层面的工具系统级工具和框架级工具。系统级工具nvidia-smi与pynvml这是与NVIDIA GPU硬件驱动直接交互的底层工具。nvidia-smi是NVIDIA提供的命令行工具功能强大信息全面。但在Python中直接调用命令行并解析输出比较繁琐。因此NVIDIA官方提供了它的Python封装pynvml库通常通过pip install nvidia-ml-py安装。pynvml允许你以编程方式获取几乎nvidia-smi能提供的所有信息是进行GPU探测和状态监控的黄金标准。它的优势在于通用、底层、信息准确不依赖于任何深度学习框架。框架级工具PyTorch 与 TensorFlow这是实际执行计算的上层框架。它们都内置了GPU相关的操作接口。PyTorch使用torch.cuda模块。它非常直观torch.cuda.is_available()检查CUDA是否可用torch.cuda.device_count()获取GPU数量torch.cuda.set_device(device_id)用于设置当前线程使用的GPU。TensorFlow在TF 2.x中通常使用tf.config.list_physical_devices(GPU)来列出GPU并通过tf.config.set_visible_devices()和tf.config.set_logical_device_configuration()来管理GPU可见性和内存分配策略。选型逻辑与搭配建议一个健壮的方案通常是“pynvml探测 深度学习框架绑定”的组合。为什么不用框架自带的查看功能框架自带的功能如torch.cuda.device_count()通常只能告诉你“有几块卡”但无法得知每块卡的实时负载如显存使用率、GPU利用率。在共享环境中一块卡即使存在也可能已被其他任务占满显存此时盲目绑定上去会导致你的任务因OOM内存溢出而失败。pynvml提供的实时状态信息是做出明智选择的关键。分工明确pynvml负责侦察和决策哪块卡最合适深度学习框架负责执行和绑定让计算任务跑在指定的卡上。注意本文主要围绕最常见的NVIDIA GPU和CUDA生态进行讲解。对于AMD GPUROCm生态或苹果M系列芯片Metal核心思路相通但具体工具和API不同需要参考对应官方文档。3. 核心细节解析与实操要点3.1 使用 pynvml 进行深度侦察仅仅知道有几块GPU是不够的。我们需要像系统管理员一样查看每块GPU的“健康状态”和“工作负荷”。pynvml库就是我们的仪表盘。首先确保已安装驱动和该库pip install nvidia-ml-py下面是一个功能增强的GPU侦察脚本它提供了比简单计数更有价值的信息import pynvml def get_gpu_status(): 获取所有GPU的详细状态信息。 返回一个列表其中每个元素是一个字典包含单块GPU的详细信息。 pynvml.nvmlInit() device_count pynvml.nvmlDeviceGetCount() gpu_list [] for i in range(device_count): handle pynvml.nvmlDeviceGetHandleByIndex(i) # 获取GPU名称 name pynvml.nvmlDeviceGetName(handle).decode(utf-8) # 获取内存信息单位字节 mem_info pynvml.nvmlDeviceGetMemoryInfo(handle) total_mem mem_info.total / (1024**3) # 转换为GB used_mem mem_info.used / (1024**3) free_mem mem_info.free / (1024**3) # 获取GPU利用率计算单元负载百分比 utilization pynvml.nvmlDeviceGetUtilizationRates(handle) gpu_util utilization.gpu # GPU计算利用率 # 注意memory_utilization 在 pynvml 中不是所有驱动都支持这里用内存使用率代替 mem_util (used_mem / total_mem) * 100 # 获取温度可选 try: temperature pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU) except pynvml.NVML_ERROR_NOT_SUPPORTED: temperature None # 获取当前运行进程信息对于判断“谁在用”至关重要 try: processes pynvml.nvmlDeviceGetComputeRunningProcesses(handle) process_list [{pid: p.pid, used_mem_MB: p.usedGpuMemory / (1024**2)} for p in processes] except pynvml.NVML_ERROR_NOT_SUPPORTED: process_list [] gpu_info { index: i, name: name, total_mem_GB: round(total_mem, 2), used_mem_GB: round(used_mem, 2), free_mem_GB: round(free_mem, 2), gpu_util_percent: gpu_util, mem_util_percent: round(mem_util, 1), temperature_C: temperature, processes: process_list, is_available: free_mem 1.0 and gpu_util 50 # 一个简单的可用性判断逻辑 } gpu_list.append(gpu_info) pynvml.nvmlShutdown() return gpu_list if __name__ __main__: gpus get_gpu_status() for gpu in gpus: print(fGPU {gpu[index]}: {gpu[name]}) print(f 显存: {gpu[used_mem_GB]:.1f} / {gpu[total_mem_GB]:.1f} GB (占用率 {gpu[mem_util_percent]}%)) print(f 计算利用率: {gpu[gpu_util_percent]}%) print(f 可用性(自定义): {可用 if gpu[is_available] else 繁忙/显存不足}) if gpu[processes]: print(f 运行中的进程: {gpu[processes]}) print(- * 40)关键点解析与实操心得nvmlInit()与nvmlShutdown()这是pynvml的标准起手式和收尾式必须成对出现确保资源正确释放。内存单位转换nvmlDeviceGetMemoryInfo返回的单位是字节。为了人类可读我们通常转换为GB除以1024^3。在判断可用性时free_mem 1.0这个条件即剩余显存大于1GB是一个经验阈值你可以根据你的模型大小调整。一个只有几百MB显存的卡很可能跑不动稍大一点的模型。利用率 vs 可用性gpu_util显示的是GPU计算核心的繁忙程度而mem_util显示的是显存占用比例。一块卡可能计算利用率很低0%但显存被一个休眠的进程占满了mem_util 90%它对你来说依然是“不可用”的。因此判断可用性必须同时考虑显存和计算负载。进程信息nvmlDeviceGetComputeRunningProcesses能告诉你当前GPU上正在运行哪些进程及其占用的显存。这在共享服务器上排查冲突、确认自己的任务是否成功绑定到正确GPU时非常有用。错误处理像获取温度这样的功能在某些显卡或驱动上可能不被支持代码中用try-except进行了包裹防止程序崩溃。3.2 制定GPU选择策略拿到所有GPU的详细信息后我们需要一个策略来做出选择。策略取决于你的目标目标A找一块最空闲的卡跑单个任务。策略优先选择free_mem最大且gpu_util最小的卡。可以给两者赋予权重进行打分。例如score free_mem_GB * 0.7 - gpu_util_percent * 0.3选择分数最高的。目标B为需要大量显存的大模型任务找卡。策略设定一个最低显存要求例如需要10GB。过滤出free_mem_GB 10的卡再从这些卡中选择gpu_util最低的。目标C手动指定或使用环境变量。策略这是最常见和灵活的方式。通过命令行参数、配置文件或环境变量如CUDA_VISIBLE_DEVICES传入一个或多个GPU索引程序直接使用。这允许用户在运行前根据侦察结果手动决策。下面是一个实现“目标A”的简单策略函数def select_most_idle_gpu(gpu_info_list): 根据自定义的闲散度评分选择最空闲的GPU。 评分规则空闲显存(GB)权重0.7空闲计算能力(100-利用率)权重0.3。 if not gpu_info_list: return None best_gpu None best_score -float(inf) for gpu in gpu_info_list: # 计算闲散度分数 free_mem_score gpu[free_mem_GB] free_util_score 100 - gpu[gpu_util_percent] # 计算空闲率 # 加权总分 score free_mem_score * 0.7 free_util_score * 0.3 if score best_score: best_score score best_gpu gpu print(f策略选择结果GPU {best_gpu[index]}得分 {best_score:.2f}) return best_gpu[index]4. 实操过程在PyTorch和TensorFlow中绑定GPU侦察完成策略选定接下来就是让框架“听话”地在指定GPU上工作。4.1 PyTorch 中的GPU指定PyTorch的设计非常Pythonic指定GPU主要有两种方式方式一使用torch.cuda.set_device(推荐用于单卡)这是在进程/线程层面设置当前上下文使用的默认GPU。import torch # 假设通过上面的策略我们决定使用 index 1 的GPU selected_gpu_index 1 # 方法1设置当前设备 torch.cuda.set_device(selected_gpu_index) # 之后创建的所有Tensor和模型如果没有特别指定device都会放在这个GPU上 device torch.device(cuda if torch.cuda.is_available() else cpu) # 此时device 对应的是 cuda:1 (如果selected_gpu_index1) model MyModel().to(device) data data.to(device)方式二在.to(device)中显式指定更精确的做法是在移动每个模型或张量时直接指定目标设备。import torch # 明确指定设备字符串 device torch.device(fcuda:{selected_gpu_index}) model MyModel().to(device) data data.to(device) # 或者在多个GPU时可以分别处理不同的部分模型并行 # model_part1.to(cuda:0) # model_part2.to(cuda:1)方式三通过环境变量CUDA_VISIBLE_DEVICES(最底层)这是最底层、最彻底的方法。它不是在Python代码中设置而是在启动Python解释器之前通过操作系统环境变量来控制PyTorch以及任何基于CUDA的程序“看到”哪些GPU。# 在Linux/macOS终端或Windows CMD中运行脚本前设置 export CUDA_VISIBLE_DEVICES1,3 # 只让程序看到物理GPU 1和3它们会被重新编号为0和1 python your_script.py # 在Python脚本内部设置需在import torch之前 import os os.environ[CUDA_VISIBLE_DEVICES] 1,3 import torch # 必须在设置环境变量之后import重要提示CUDA_VISIBLE_DEVICES的优先级最高。如果你在代码里torch.cuda.set_device(0)但环境变量设置的是CUDA_VISIBLE_DEVICES2那么你实际操作的物理GPU是2号卡它被重映射为逻辑0号卡。混用时务必小心。4.2 TensorFlow 2.x 中的GPU指定TensorFlow 2.x的GPU管理API更倾向于显式地设置“可见设备列表”。方式一使用tf.config.set_visible_devicesimport tensorflow as tf # 获取所有物理GPU设备 gpus tf.config.list_physical_devices(GPU) if gpus: # 例如我们只想使用第0号和第2号物理GPU try: # 设置只有哪些GPU对TF可见 tf.config.set_visible_devices([gpus[0], gpus[2]], GPU) # 可选设置内存动态增长避免一开始占满所有显存 for gpu in [gpus[0], gpus[2]]: tf.config.experimental.set_memory_growth(gpu, True) print(fTensorFlow将使用物理GPU: 0 和 2) except RuntimeError as e: # 虚拟设备必须在运行时启动前设置 print(e)方式二也支持环境变量CUDA_VISIBLE_DEVICES和PyTorch一样TensorFlow也完全遵从CUDA_VISIBLE_DEVICES环境变量。这是跨框架的统一方法。方式三在创建策略时指定用于分布式训练对于多卡训练你可以在创建MirroredStrategy等分布策略时指定设备。# 创建一个只使用GPU 0和1的分布式策略 strategy tf.distribute.MirroredStrategy(devices[/gpu:0, /gpu:1]) with strategy.scope(): # 在这里构建和编译你的模型 model ...4.3 一个完整的自动化脚本示例将侦察、选择、绑定三个步骤串联起来形成一个完整的自动化工作流import pynvml import torch import argparse def get_gpu_info(): 获取GPU信息列表 # ... (使用前面定义的 get_gpu_status 函数) ... pass def auto_select_gpu(gpu_list, min_free_mem_gb4.0, max_gpu_util30.0): 自动选择GPU。 策略找到第一块满足【空闲显存 min_free_mem_gb】且【GPU利用率 max_gpu_util】的卡。 for gpu in gpu_list: if gpu[free_mem_GB] min_free_mem_gb and gpu[gpu_util_percent] max_gpu_util: return gpu[index] print(警告未找到符合条件空闲显存{}GB且利用率{}%的GPU。将使用策略选择最空闲的。.format(min_free_mem_gb, max_gpu_util)) # 降级策略选择最空闲的 return select_most_idle_gpu(gpu_list) def main(): parser argparse.ArgumentParser() parser.add_argument(--gpu, typeint, defaultNone, help手动指定GPU索引例如 0, 1。不指定则自动选择。) parser.add_argument(--min-free-mem, typefloat, default4.0, help自动选择时要求的最小空闲显存(GB)) parser.add_argument(--max-gpu-util, typefloat, default30.0, help自动选择时允许的最大GPU利用率(%)) args parser.parse_args() # 1. 侦察 gpu_list get_gpu_info() print( 系统GPU状态 ) for g in gpu_list: print(fGPU {g[index]}: {g[name]}, 空闲显存 {g[free_mem_GB]:.1f}GB, 利用率 {g[gpu_util_percent]}%) # 2. 选择 selected_index args.gpu if selected_index is None: selected_index auto_select_gpu(gpu_list, args.min_free_mem, args.max_gpu_util) print(f\n[自动选择] 将使用 GPU {selected_index}) else: # 检查手动指定的GPU是否有效 if selected_index 0 or selected_index len(gpu_list): print(f错误指定的GPU索引 {selected_index} 无效。) return print(f\n[手动指定] 将使用 GPU {selected_index}) # 3. 绑定 (以PyTorch为例) if torch.cuda.is_available(): # 方法A设置当前设备 torch.cuda.set_device(selected_index) current_device torch.cuda.current_device() print(fPyTorch当前设备已设置为: cuda:{current_device}) # 验证 test_tensor torch.tensor([1.0, 2.0, 3.0]).cuda() print(f测试张量所在设备: {test_tensor.device}) else: print(CUDA不可用将使用CPU。) # 4. 你的训练/推理代码从这里开始... # model MyModel().cuda() # ... if __name__ __main__: main()这个脚本提供了灵活的接口可以通过--gpu手动指定也可以通过--min-free-mem和--max-gpu-util参数调整自动选择策略。在程序开始时打印出清晰的GPU状态和最终选择结果便于记录和调试。5. 常见问题与排查技巧实录在实际操作中你肯定会遇到各种“坑”。下面是我总结的一些典型问题及其解决方法。5.1 问题一代码指定了GPU但任务仍然跑在了其他卡上现象明明用了torch.cuda.set_device(0)但通过nvidia-smi查看发现进程跑在了GPU 2上。原因与排查环境变量覆盖这是最常见的原因。检查是否在运行脚本时或脚本开头设置了CUDA_VISIBLE_DEVICES。例如CUDA_VISIBLE_DEVICES2 python script.py会让程序只看到物理GPU 2它在程序内被识别为cuda:0。此时你在代码里设置set_device(0)操作的正是物理GPU 2。排查在代码中打印os.environ.get(CUDA_VISIBLE_DEVICES)和torch.cuda.current_device()进行对比。代码中存在未指定设备的.cuda()调用在PyTorch的旧代码或某些示例中会直接使用.cuda()将模型或数据放到GPU上。如果不带参数.cuda()会默认放到“当前”GPU而“当前”GPU可能被其他地方的set_device改变导致混乱。解决永远使用.to(device)方法并显式传入device对象。这是最清晰、最不容易出错的方式。# 推荐做法 device torch.device(fcuda:{selected_gpu_index}) model.to(device) data.to(device) # 避免使用除非你非常清楚上下文 model.cuda() # 不推荐 model.cuda(1) # 可以但不如.to(device)清晰多进程/多线程问题如果你使用了Python的multiprocessing或DataLoader的多进程加载num_workers 0每个子进程都会继承父进程的环境但可能需要重新设置设备。在子进程函数内部最好也显式地设置一次torch.cuda.set_device。5.2 问题二显存溢出CUDA out of memory现象程序开始运行不久就报错RuntimeError: CUDA out of memory。排查步骤确认目标GPU的剩余显存运行前用我们上面的侦察脚本确认free_mem_GB是否真的足够你的模型和数据。一个常见的错觉是GPU利用率低就等于显存空闲。显存是被事先分配和占用的即使计算核心空闲显存也可能已满。检查是否有其他进程占用使用侦察脚本中的processes字段或直接在终端运行nvidia-smi查看目标GPU上是否有其他进程如别人的训练任务、僵尸进程占用了大量显存。排查代码中的显存泄漏张量累积在训练循环中是否将损失、日志等张量不断追加到列表中这会导致计算图不断增长显存无法释放。正确的做法是使用.item()将标量损失转换为Python数字。# 错误做法 losses [] for data, target in dataloader: output model(data) loss criterion(output, target) losses.append(loss) # loss是张量会保留计算图 ... # 正确做法 total_loss 0 for data, target in dataloader: output model(data) loss criterion(output, target) total_loss loss.item() # 使用.item()获取标量值 loss.backward() ...未及时释放无用变量在循环中创建的大中间变量如果不再需要可以手动将其从GPU移出并删除。intermediate some_heavy_operation(data) # 一个大张量 result process(intermediate) # 不再需要intermediate了 intermediate intermediate.cpu() # 移到CPU del intermediate # 删除引用 torch.cuda.empty_cache() # 可选通知CUDA释放缓存但不要频繁调用验证/测试阶段忘记torch.no_grad()在不需要计算梯度的前向传播阶段如验证、测试务必使用with torch.no_grad():上下文管理器这可以避免保存中间变量用于反向传播极大节省显存。调整批次大小Batch Size这是最直接的杠杆。如果显存不足首先尝试减小batch_size。5.3 问题三多卡训练时负载不均衡或速度没提升现象使用了DataParallel或DistributedDataParallel但发现只有一张卡在忙或者速度比单卡还慢。排查与解决确认数据并行正确性对于DataParallel它自动将数据拆分到各卡。确保你的模型确实在DataParallel包装下。model MyModel() if torch.cuda.device_count() 1: print(f使用 {torch.cuda.device_count()} 块GPU进行数据并行) model torch.nn.DataParallel(model) model.to(device)但请注意DataParallel是单进程多线程的在有些情况下效率不如DistributedDataParallel。检查数据加载瓶颈多卡训练时数据加载可能成为瓶颈。确保你的DataLoader使用了足够数量的工作进程num_workers并且数据预处理不是太重。可以使用torch.utils.data.DataLoader的pin_memoryTrue选项加速数据从CPU到GPU的传输。考虑使用DistributedDataParallel (DDP)对于真正的多机多卡或单机多卡大规模训练DistributedDataParallel是更推荐的方式。它采用多进程方式每个进程对应一块GPU避免了DataParallel的全局解释器锁GIL限制和单进程通信开销。虽然设置稍复杂但通常能获得更好的扩展性。监控每张卡的利用率在训练过程中另开一个终端使用watch -n 1 nvidia-smi命令实时观察每张GPU的利用率和显存占用。如果只有一张卡利用率高可能是数据没有正确分发。5.4 环境配置检查清单在遇到任何GPU相关问题时可以按以下清单快速排查基础环境CUDA驱动与运行时版本匹配命令行运行nvidia-smi查看右上角的“CUDA Version”这是驱动支持的最高CUDA运行时版本。在Python中运行torch.version.cuda(PyTorch) 或tf.sysconfig.get_build_info()[cuda_version](TensorFlow)查看框架实际编译使用的CUDA运行时版本。要求框架所需的CUDA运行时版本必须小于等于驱动支持的版本。例如驱动支持CUDA 12.2你可以安装基于CUDA 11.8或12.1编译的PyTorch但不能安装需要CUDA 12.3的。PyTorch/TensorFlow安装是否正确PyTorch:print(torch.cuda.is_available())必须返回True。TensorFlow 2.x:print(tf.config.list_physical_devices(GPU))应该能列出GPU设备。如果返回False或空列表但nvidia-smi正常大概率是框架的CUDA版本与系统环境不匹配需要重新安装对应版本的框架。权限问题Linux系统常见当前用户是否有权限访问GPU设备通常是/dev/nvidia*文件。在某些Docker容器或严格配置的服务器上可能需要将用户加入video或render组。掌握查看和指定GPU的方法是高效利用计算资源的第一步。它让你从被动的“能用就行”转变为主动的“精准控制”。从基础的nvidia-smi和torch.cuda.set_device到结合pynvml进行智能选择再到理解环境变量CUDA_VISIBLE_DEVICES的底层逻辑每一步都让你对计算环境有更深的掌控力。记住清晰的策略、显式的设备指定多用.to(device)和实时的状态监控nvidia-smi是避免GPU相关bug的三个法宝。下次启动你的训练脚本前不妨花一分钟运行一下侦察代码或许就能避开一个长达数小时的排队等待或者莫名其妙的显存溢出错误。 Python GPU资源管理：从pynvml侦察到PyTorch/TensorFlow指定GPU实战
1. 项目概述为什么需要查看并指定GPU在深度学习、科学计算或者大规模数据处理的项目里GPU图形处理器早已不是游戏玩家的专属。它凭借其强大的并行计算能力成为了加速模型训练和复杂运算的“核动力引擎”。然而当你兴冲冲地准备跑一个大型模型时可能会遇到几个尴尬的场景服务器上有多块GPU但你的代码默认只用了一块其他几块在“围观”或者你和其他人共用一台服务器你的任务把别人的任务给挤掉了导致冲突。这时候学会用Python代码“侦察”可用的GPU并精确地“指挥”你的任务在指定的GPU上运行就从一个加分项变成了必备技能。这不仅仅是运行一个命令那么简单。它背后涉及到对计算资源的有效管理、对任务运行环境的精确控制以及对多GPU系统工作方式的理解。无论是个人开发者在一台多卡工作站上调试模型还是团队在共享的GPU集群上部署任务掌握这项技能都能让你事半功倍避免资源浪费和潜在冲突。接下来我将从一个实践者的角度带你一步步拆解这个需求从原理到实操再到避坑指南让你彻底搞懂。2. 核心需求解析与方案选型2.1 需求拆解我们到底要做什么看似简单的“查看并指定GPU”其实可以拆解为三个递进的核心动作探测Probe程序启动时自动识别当前系统中有哪些GPU设备以及它们的关键状态如是否空闲、显存占用、计算负载等。这相当于给你的代码装上了“雷达”。选择Select基于探测到的信息结合你的任务需求例如需要大显存或者需要避开繁忙的卡制定一个选择策略决定使用哪一块或哪几块GPU。绑定Bind将你的Python进程以及进程中用到的深度学习框架如PyTorch、TensorFlow的计算任务强制绑定到你所选定的GPU设备上确保任务不会“跑偏”。2.2 主流工具选型为什么是它们在Python生态中完成这些动作主要依赖两个层面的工具系统级工具和框架级工具。系统级工具nvidia-smi与pynvml这是与NVIDIA GPU硬件驱动直接交互的底层工具。nvidia-smi是NVIDIA提供的命令行工具功能强大信息全面。但在Python中直接调用命令行并解析输出比较繁琐。因此NVIDIA官方提供了它的Python封装pynvml库通常通过pip install nvidia-ml-py安装。pynvml允许你以编程方式获取几乎nvidia-smi能提供的所有信息是进行GPU探测和状态监控的黄金标准。它的优势在于通用、底层、信息准确不依赖于任何深度学习框架。框架级工具PyTorch 与 TensorFlow这是实际执行计算的上层框架。它们都内置了GPU相关的操作接口。PyTorch使用torch.cuda模块。它非常直观torch.cuda.is_available()检查CUDA是否可用torch.cuda.device_count()获取GPU数量torch.cuda.set_device(device_id)用于设置当前线程使用的GPU。TensorFlow在TF 2.x中通常使用tf.config.list_physical_devices(GPU)来列出GPU并通过tf.config.set_visible_devices()和tf.config.set_logical_device_configuration()来管理GPU可见性和内存分配策略。选型逻辑与搭配建议一个健壮的方案通常是“pynvml探测 深度学习框架绑定”的组合。为什么不用框架自带的查看功能框架自带的功能如torch.cuda.device_count()通常只能告诉你“有几块卡”但无法得知每块卡的实时负载如显存使用率、GPU利用率。在共享环境中一块卡即使存在也可能已被其他任务占满显存此时盲目绑定上去会导致你的任务因OOM内存溢出而失败。pynvml提供的实时状态信息是做出明智选择的关键。分工明确pynvml负责侦察和决策哪块卡最合适深度学习框架负责执行和绑定让计算任务跑在指定的卡上。注意本文主要围绕最常见的NVIDIA GPU和CUDA生态进行讲解。对于AMD GPUROCm生态或苹果M系列芯片Metal核心思路相通但具体工具和API不同需要参考对应官方文档。3. 核心细节解析与实操要点3.1 使用 pynvml 进行深度侦察仅仅知道有几块GPU是不够的。我们需要像系统管理员一样查看每块GPU的“健康状态”和“工作负荷”。pynvml库就是我们的仪表盘。首先确保已安装驱动和该库pip install nvidia-ml-py下面是一个功能增强的GPU侦察脚本它提供了比简单计数更有价值的信息import pynvml def get_gpu_status(): 获取所有GPU的详细状态信息。 返回一个列表其中每个元素是一个字典包含单块GPU的详细信息。 pynvml.nvmlInit() device_count pynvml.nvmlDeviceGetCount() gpu_list [] for i in range(device_count): handle pynvml.nvmlDeviceGetHandleByIndex(i) # 获取GPU名称 name pynvml.nvmlDeviceGetName(handle).decode(utf-8) # 获取内存信息单位字节 mem_info pynvml.nvmlDeviceGetMemoryInfo(handle) total_mem mem_info.total / (1024**3) # 转换为GB used_mem mem_info.used / (1024**3) free_mem mem_info.free / (1024**3) # 获取GPU利用率计算单元负载百分比 utilization pynvml.nvmlDeviceGetUtilizationRates(handle) gpu_util utilization.gpu # GPU计算利用率 # 注意memory_utilization 在 pynvml 中不是所有驱动都支持这里用内存使用率代替 mem_util (used_mem / total_mem) * 100 # 获取温度可选 try: temperature pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU) except pynvml.NVML_ERROR_NOT_SUPPORTED: temperature None # 获取当前运行进程信息对于判断“谁在用”至关重要 try: processes pynvml.nvmlDeviceGetComputeRunningProcesses(handle) process_list [{pid: p.pid, used_mem_MB: p.usedGpuMemory / (1024**2)} for p in processes] except pynvml.NVML_ERROR_NOT_SUPPORTED: process_list [] gpu_info { index: i, name: name, total_mem_GB: round(total_mem, 2), used_mem_GB: round(used_mem, 2), free_mem_GB: round(free_mem, 2), gpu_util_percent: gpu_util, mem_util_percent: round(mem_util, 1), temperature_C: temperature, processes: process_list, is_available: free_mem 1.0 and gpu_util 50 # 一个简单的可用性判断逻辑 } gpu_list.append(gpu_info) pynvml.nvmlShutdown() return gpu_list if __name__ __main__: gpus get_gpu_status() for gpu in gpus: print(fGPU {gpu[index]}: {gpu[name]}) print(f 显存: {gpu[used_mem_GB]:.1f} / {gpu[total_mem_GB]:.1f} GB (占用率 {gpu[mem_util_percent]}%)) print(f 计算利用率: {gpu[gpu_util_percent]}%) print(f 可用性(自定义): {可用 if gpu[is_available] else 繁忙/显存不足}) if gpu[processes]: print(f 运行中的进程: {gpu[processes]}) print(- * 40)关键点解析与实操心得nvmlInit()与nvmlShutdown()这是pynvml的标准起手式和收尾式必须成对出现确保资源正确释放。内存单位转换nvmlDeviceGetMemoryInfo返回的单位是字节。为了人类可读我们通常转换为GB除以1024^3。在判断可用性时free_mem 1.0这个条件即剩余显存大于1GB是一个经验阈值你可以根据你的模型大小调整。一个只有几百MB显存的卡很可能跑不动稍大一点的模型。利用率 vs 可用性gpu_util显示的是GPU计算核心的繁忙程度而mem_util显示的是显存占用比例。一块卡可能计算利用率很低0%但显存被一个休眠的进程占满了mem_util 90%它对你来说依然是“不可用”的。因此判断可用性必须同时考虑显存和计算负载。进程信息nvmlDeviceGetComputeRunningProcesses能告诉你当前GPU上正在运行哪些进程及其占用的显存。这在共享服务器上排查冲突、确认自己的任务是否成功绑定到正确GPU时非常有用。错误处理像获取温度这样的功能在某些显卡或驱动上可能不被支持代码中用try-except进行了包裹防止程序崩溃。3.2 制定GPU选择策略拿到所有GPU的详细信息后我们需要一个策略来做出选择。策略取决于你的目标目标A找一块最空闲的卡跑单个任务。策略优先选择free_mem最大且gpu_util最小的卡。可以给两者赋予权重进行打分。例如score free_mem_GB * 0.7 - gpu_util_percent * 0.3选择分数最高的。目标B为需要大量显存的大模型任务找卡。策略设定一个最低显存要求例如需要10GB。过滤出free_mem_GB 10的卡再从这些卡中选择gpu_util最低的。目标C手动指定或使用环境变量。策略这是最常见和灵活的方式。通过命令行参数、配置文件或环境变量如CUDA_VISIBLE_DEVICES传入一个或多个GPU索引程序直接使用。这允许用户在运行前根据侦察结果手动决策。下面是一个实现“目标A”的简单策略函数def select_most_idle_gpu(gpu_info_list): 根据自定义的闲散度评分选择最空闲的GPU。 评分规则空闲显存(GB)权重0.7空闲计算能力(100-利用率)权重0.3。 if not gpu_info_list: return None best_gpu None best_score -float(inf) for gpu in gpu_info_list: # 计算闲散度分数 free_mem_score gpu[free_mem_GB] free_util_score 100 - gpu[gpu_util_percent] # 计算空闲率 # 加权总分 score free_mem_score * 0.7 free_util_score * 0.3 if score best_score: best_score score best_gpu gpu print(f策略选择结果GPU {best_gpu[index]}得分 {best_score:.2f}) return best_gpu[index]4. 实操过程在PyTorch和TensorFlow中绑定GPU侦察完成策略选定接下来就是让框架“听话”地在指定GPU上工作。4.1 PyTorch 中的GPU指定PyTorch的设计非常Pythonic指定GPU主要有两种方式方式一使用torch.cuda.set_device(推荐用于单卡)这是在进程/线程层面设置当前上下文使用的默认GPU。import torch # 假设通过上面的策略我们决定使用 index 1 的GPU selected_gpu_index 1 # 方法1设置当前设备 torch.cuda.set_device(selected_gpu_index) # 之后创建的所有Tensor和模型如果没有特别指定device都会放在这个GPU上 device torch.device(cuda if torch.cuda.is_available() else cpu) # 此时device 对应的是 cuda:1 (如果selected_gpu_index1) model MyModel().to(device) data data.to(device)方式二在.to(device)中显式指定更精确的做法是在移动每个模型或张量时直接指定目标设备。import torch # 明确指定设备字符串 device torch.device(fcuda:{selected_gpu_index}) model MyModel().to(device) data data.to(device) # 或者在多个GPU时可以分别处理不同的部分模型并行 # model_part1.to(cuda:0) # model_part2.to(cuda:1)方式三通过环境变量CUDA_VISIBLE_DEVICES(最底层)这是最底层、最彻底的方法。它不是在Python代码中设置而是在启动Python解释器之前通过操作系统环境变量来控制PyTorch以及任何基于CUDA的程序“看到”哪些GPU。# 在Linux/macOS终端或Windows CMD中运行脚本前设置 export CUDA_VISIBLE_DEVICES1,3 # 只让程序看到物理GPU 1和3它们会被重新编号为0和1 python your_script.py # 在Python脚本内部设置需在import torch之前 import os os.environ[CUDA_VISIBLE_DEVICES] 1,3 import torch # 必须在设置环境变量之后import重要提示CUDA_VISIBLE_DEVICES的优先级最高。如果你在代码里torch.cuda.set_device(0)但环境变量设置的是CUDA_VISIBLE_DEVICES2那么你实际操作的物理GPU是2号卡它被重映射为逻辑0号卡。混用时务必小心。4.2 TensorFlow 2.x 中的GPU指定TensorFlow 2.x的GPU管理API更倾向于显式地设置“可见设备列表”。方式一使用tf.config.set_visible_devicesimport tensorflow as tf # 获取所有物理GPU设备 gpus tf.config.list_physical_devices(GPU) if gpus: # 例如我们只想使用第0号和第2号物理GPU try: # 设置只有哪些GPU对TF可见 tf.config.set_visible_devices([gpus[0], gpus[2]], GPU) # 可选设置内存动态增长避免一开始占满所有显存 for gpu in [gpus[0], gpus[2]]: tf.config.experimental.set_memory_growth(gpu, True) print(fTensorFlow将使用物理GPU: 0 和 2) except RuntimeError as e: # 虚拟设备必须在运行时启动前设置 print(e)方式二也支持环境变量CUDA_VISIBLE_DEVICES和PyTorch一样TensorFlow也完全遵从CUDA_VISIBLE_DEVICES环境变量。这是跨框架的统一方法。方式三在创建策略时指定用于分布式训练对于多卡训练你可以在创建MirroredStrategy等分布策略时指定设备。# 创建一个只使用GPU 0和1的分布式策略 strategy tf.distribute.MirroredStrategy(devices[/gpu:0, /gpu:1]) with strategy.scope(): # 在这里构建和编译你的模型 model ...4.3 一个完整的自动化脚本示例将侦察、选择、绑定三个步骤串联起来形成一个完整的自动化工作流import pynvml import torch import argparse def get_gpu_info(): 获取GPU信息列表 # ... (使用前面定义的 get_gpu_status 函数) ... pass def auto_select_gpu(gpu_list, min_free_mem_gb4.0, max_gpu_util30.0): 自动选择GPU。 策略找到第一块满足【空闲显存 min_free_mem_gb】且【GPU利用率 max_gpu_util】的卡。 for gpu in gpu_list: if gpu[free_mem_GB] min_free_mem_gb and gpu[gpu_util_percent] max_gpu_util: return gpu[index] print(警告未找到符合条件空闲显存{}GB且利用率{}%的GPU。将使用策略选择最空闲的。.format(min_free_mem_gb, max_gpu_util)) # 降级策略选择最空闲的 return select_most_idle_gpu(gpu_list) def main(): parser argparse.ArgumentParser() parser.add_argument(--gpu, typeint, defaultNone, help手动指定GPU索引例如 0, 1。不指定则自动选择。) parser.add_argument(--min-free-mem, typefloat, default4.0, help自动选择时要求的最小空闲显存(GB)) parser.add_argument(--max-gpu-util, typefloat, default30.0, help自动选择时允许的最大GPU利用率(%)) args parser.parse_args() # 1. 侦察 gpu_list get_gpu_info() print( 系统GPU状态 ) for g in gpu_list: print(fGPU {g[index]}: {g[name]}, 空闲显存 {g[free_mem_GB]:.1f}GB, 利用率 {g[gpu_util_percent]}%) # 2. 选择 selected_index args.gpu if selected_index is None: selected_index auto_select_gpu(gpu_list, args.min_free_mem, args.max_gpu_util) print(f\n[自动选择] 将使用 GPU {selected_index}) else: # 检查手动指定的GPU是否有效 if selected_index 0 or selected_index len(gpu_list): print(f错误指定的GPU索引 {selected_index} 无效。) return print(f\n[手动指定] 将使用 GPU {selected_index}) # 3. 绑定 (以PyTorch为例) if torch.cuda.is_available(): # 方法A设置当前设备 torch.cuda.set_device(selected_index) current_device torch.cuda.current_device() print(fPyTorch当前设备已设置为: cuda:{current_device}) # 验证 test_tensor torch.tensor([1.0, 2.0, 3.0]).cuda() print(f测试张量所在设备: {test_tensor.device}) else: print(CUDA不可用将使用CPU。) # 4. 你的训练/推理代码从这里开始... # model MyModel().cuda() # ... if __name__ __main__: main()这个脚本提供了灵活的接口可以通过--gpu手动指定也可以通过--min-free-mem和--max-gpu-util参数调整自动选择策略。在程序开始时打印出清晰的GPU状态和最终选择结果便于记录和调试。5. 常见问题与排查技巧实录在实际操作中你肯定会遇到各种“坑”。下面是我总结的一些典型问题及其解决方法。5.1 问题一代码指定了GPU但任务仍然跑在了其他卡上现象明明用了torch.cuda.set_device(0)但通过nvidia-smi查看发现进程跑在了GPU 2上。原因与排查环境变量覆盖这是最常见的原因。检查是否在运行脚本时或脚本开头设置了CUDA_VISIBLE_DEVICES。例如CUDA_VISIBLE_DEVICES2 python script.py会让程序只看到物理GPU 2它在程序内被识别为cuda:0。此时你在代码里设置set_device(0)操作的正是物理GPU 2。排查在代码中打印os.environ.get(CUDA_VISIBLE_DEVICES)和torch.cuda.current_device()进行对比。代码中存在未指定设备的.cuda()调用在PyTorch的旧代码或某些示例中会直接使用.cuda()将模型或数据放到GPU上。如果不带参数.cuda()会默认放到“当前”GPU而“当前”GPU可能被其他地方的set_device改变导致混乱。解决永远使用.to(device)方法并显式传入device对象。这是最清晰、最不容易出错的方式。# 推荐做法 device torch.device(fcuda:{selected_gpu_index}) model.to(device) data.to(device) # 避免使用除非你非常清楚上下文 model.cuda() # 不推荐 model.cuda(1) # 可以但不如.to(device)清晰多进程/多线程问题如果你使用了Python的multiprocessing或DataLoader的多进程加载num_workers 0每个子进程都会继承父进程的环境但可能需要重新设置设备。在子进程函数内部最好也显式地设置一次torch.cuda.set_device。5.2 问题二显存溢出CUDA out of memory现象程序开始运行不久就报错RuntimeError: CUDA out of memory。排查步骤确认目标GPU的剩余显存运行前用我们上面的侦察脚本确认free_mem_GB是否真的足够你的模型和数据。一个常见的错觉是GPU利用率低就等于显存空闲。显存是被事先分配和占用的即使计算核心空闲显存也可能已满。检查是否有其他进程占用使用侦察脚本中的processes字段或直接在终端运行nvidia-smi查看目标GPU上是否有其他进程如别人的训练任务、僵尸进程占用了大量显存。排查代码中的显存泄漏张量累积在训练循环中是否将损失、日志等张量不断追加到列表中这会导致计算图不断增长显存无法释放。正确的做法是使用.item()将标量损失转换为Python数字。# 错误做法 losses [] for data, target in dataloader: output model(data) loss criterion(output, target) losses.append(loss) # loss是张量会保留计算图 ... # 正确做法 total_loss 0 for data, target in dataloader: output model(data) loss criterion(output, target) total_loss loss.item() # 使用.item()获取标量值 loss.backward() ...未及时释放无用变量在循环中创建的大中间变量如果不再需要可以手动将其从GPU移出并删除。intermediate some_heavy_operation(data) # 一个大张量 result process(intermediate) # 不再需要intermediate了 intermediate intermediate.cpu() # 移到CPU del intermediate # 删除引用 torch.cuda.empty_cache() # 可选通知CUDA释放缓存但不要频繁调用验证/测试阶段忘记torch.no_grad()在不需要计算梯度的前向传播阶段如验证、测试务必使用with torch.no_grad():上下文管理器这可以避免保存中间变量用于反向传播极大节省显存。调整批次大小Batch Size这是最直接的杠杆。如果显存不足首先尝试减小batch_size。5.3 问题三多卡训练时负载不均衡或速度没提升现象使用了DataParallel或DistributedDataParallel但发现只有一张卡在忙或者速度比单卡还慢。排查与解决确认数据并行正确性对于DataParallel它自动将数据拆分到各卡。确保你的模型确实在DataParallel包装下。model MyModel() if torch.cuda.device_count() 1: print(f使用 {torch.cuda.device_count()} 块GPU进行数据并行) model torch.nn.DataParallel(model) model.to(device)但请注意DataParallel是单进程多线程的在有些情况下效率不如DistributedDataParallel。检查数据加载瓶颈多卡训练时数据加载可能成为瓶颈。确保你的DataLoader使用了足够数量的工作进程num_workers并且数据预处理不是太重。可以使用torch.utils.data.DataLoader的pin_memoryTrue选项加速数据从CPU到GPU的传输。考虑使用DistributedDataParallel (DDP)对于真正的多机多卡或单机多卡大规模训练DistributedDataParallel是更推荐的方式。它采用多进程方式每个进程对应一块GPU避免了DataParallel的全局解释器锁GIL限制和单进程通信开销。虽然设置稍复杂但通常能获得更好的扩展性。监控每张卡的利用率在训练过程中另开一个终端使用watch -n 1 nvidia-smi命令实时观察每张GPU的利用率和显存占用。如果只有一张卡利用率高可能是数据没有正确分发。5.4 环境配置检查清单在遇到任何GPU相关问题时可以按以下清单快速排查基础环境CUDA驱动与运行时版本匹配命令行运行nvidia-smi查看右上角的“CUDA Version”这是驱动支持的最高CUDA运行时版本。在Python中运行torch.version.cuda(PyTorch) 或tf.sysconfig.get_build_info()[cuda_version](TensorFlow)查看框架实际编译使用的CUDA运行时版本。要求框架所需的CUDA运行时版本必须小于等于驱动支持的版本。例如驱动支持CUDA 12.2你可以安装基于CUDA 11.8或12.1编译的PyTorch但不能安装需要CUDA 12.3的。PyTorch/TensorFlow安装是否正确PyTorch:print(torch.cuda.is_available())必须返回True。TensorFlow 2.x:print(tf.config.list_physical_devices(GPU))应该能列出GPU设备。如果返回False或空列表但nvidia-smi正常大概率是框架的CUDA版本与系统环境不匹配需要重新安装对应版本的框架。权限问题Linux系统常见当前用户是否有权限访问GPU设备通常是/dev/nvidia*文件。在某些Docker容器或严格配置的服务器上可能需要将用户加入video或render组。掌握查看和指定GPU的方法是高效利用计算资源的第一步。它让你从被动的“能用就行”转变为主动的“精准控制”。从基础的nvidia-smi和torch.cuda.set_device到结合pynvml进行智能选择再到理解环境变量CUDA_VISIBLE_DEVICES的底层逻辑每一步都让你对计算环境有更深的掌控力。记住清晰的策略、显式的设备指定多用.to(device)和实时的状态监控nvidia-smi是避免GPU相关bug的三个法宝。下次启动你的训练脚本前不妨花一分钟运行一下侦察代码或许就能避开一个长达数小时的排队等待或者莫名其妙的显存溢出错误。