1. 项目概述为什么需要指定GPU在深度学习项目里尤其是当你手头有多张显卡时学会精确地告诉PyTorch“用哪一张卡干活”是迈向高效开发的第一步。这听起来简单但背后涉及资源管理、性能优化和团队协作的多个层面。我见过不少新手包括早期的我自己把模型往服务器上一扔就开始python train.py结果发现训练慢得出奇一查nvidia-smi宝贵的A100在摸鱼而一块老旧的Titan X却在满负荷运转。这不仅仅是浪费电费更是在浪费生命。指定GPU的核心价值在于掌控力。对于个人开发者你可能有一张用于打游戏的RTX 4090和一张用于工作的RTX 3090你肯定希望训练任务跑在3090上而不影响游戏性能。对于团队使用的服务器通常搭载着多张型号、显存各异的GPU比如4张A100 80G不同的任务对算力和显存的需求天差地别。一个需要80G显存的大语言模型微调任务必须被锁定在单张A100上而一个轻量级的图像分类任务则可能希望用数据并行分散到所有卡上以加速。如果你不指定PyTorch默认会使用CUDA_VISIBLE_DEVICES环境变量排序后的第一块卡或者torch.cuda.current_device()返回的卡这种不确定性是生产环境的大忌。更深一层指定GPU是进行多卡并行训练如DataParallel,DistributedDataParallel的基础。你必须先明确每张卡的“身份”才能有效地将模型和数据分配上去。此外在模型推理服务部署时为了服务隔离和性能稳定我们通常会将不同的模型实例绑定到不同的GPU上避免相互干扰。因此“指定GPU”这个操作是从单卡实验走向多卡协作、从本地开发走向生产部署的关键桥梁。它确保计算资源被可预测、可管理地利用是专业深度学习工作流中不可或缺的一环。2. 核心原理与环境准备2.1 CUDA、PyTorch与GPU的协作关系在动手之前我们需要理清几个核心概念之间的关系这能帮你从根本上理解后续所有操作。你可以把GPU想象成一个拥有数千个核心的超级计算车间CUDA则是这个车间的操作系统和编程模型它定义了如何向这个车间下达指令即编写核函数。而PyTorch则是一个高级的工厂管理软件。它封装了CUDA的复杂接口提供了像Tensor和nn.Module这样友好的抽象。当你执行tensor.cuda()时PyTorch这个“管理软件”就在背后调用了CUDA的API把数据从CPU内存搬运到GPU显存并安排车间的核心进行计算。那么当系统中有多块GPU时CUDA是如何标识它们的呢CUDA使用一个从0开始的整数索引来标识每块GPU。这个索引顺序通常由PCIe总线枚举顺序决定你可以通过nvidia-smi命令查看最左侧的GPU 0、GPU 1就是它们的ID。PyTorch在初始化时会与CUDA运行时建立连接并获取当前可用的GPU列表。torch.cuda.device_count()返回的就是这个数量。这里有一个至关重要的环境变量CUDA_VISIBLE_DEVICES。这个变量并不改变物理GPU的索引而是为当前进程创建一个虚拟的GPU视图。例如在一个有4块GPU0,1,2,3的服务器上如果你设置export CUDA_VISIBLE_DEVICES2,1那么在当前进程看来就只有两块GPU并且它们的索引被重新映射为0对应物理GPU 2和1对应物理GPU 1。物理GPU 0和3对这个进程来说是完全不可见的。这是最底层、最全局的GPU指定方法所有基于CUDA的应用包括PyTorch都会遵守这个视图。2.2 环境检查与工具准备在开始指定GPU之前我们必须确保环境是正确的。以下是一套完整的检查流程我习惯在任何一个新环境跑代码前都走一遍。首先打开你的终端运行nvidia-smi。这个命令的输出是你的“显卡仪表盘”。你需要关注几点Driver Version: 确保你的NVIDIA驱动版本足够新以支持你安装的CUDA版本。CUDA Version: 这里显示的是驱动所能支持的最高CUDA运行时版本不是你实际安装的PyTorch对应的CUDA版本。GPU列表: 确认所有你期望的GPU都正常识别并且没有异常的错误信息如ERR!。进程信息: 在下方表格可以看到每张GPU上正在运行的进程这能帮你判断哪张卡是空闲的。接下来在Python交互环境中用PyTorch进行验证import torch # 1. 检查PyTorch是否支持CUDA print(fCUDA available: {torch.cuda.is_available()}) # 必须为True # 2. 检查当前PyTorch构建所使用的CUDA版本 print(fPyTorch CUDA version: {torch.version.cuda}) # 3. 查看可用GPU数量 print(fNumber of GPUs: {torch.cuda.device_count()}) # 4. 查看每张GPU的名称和显存容量 for i in range(torch.cuda.device_count()): print(fGPU {i}: {torch.cuda.get_device_name(i)} | Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB)如果torch.cuda.is_available()返回False那一切免谈。常见原因有安装了CPU版本的PyTorch。请务必通过 PyTorch官网 的命令行根据你的CUDA版本选择正确的安装命令。例如对于CUDA 11.8pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118。NVIDIA驱动太旧。去官网下载更新。在虚拟环境或容器中CUDA路径未正确设置。注意torch.version.cudaPyTorch编译时用的CUDA版本与系统安装的CUDA Toolkit版本如nvcc --version显示的可以不同只要驱动支持即可。但为了兼容性建议尽量对齐。3. 训练时指定GPU的四种核心方法指定GPU不是只有一个办法根据不同的场景和需求有从粗粒度到细粒度从环境级到代码级的不同策略。理解它们的区别和适用场景能让你在复杂项目中游刃有余。3.1 方法一使用CUDA_VISIBLE_DEVICES环境变量最推荐这是我最推荐也是在生产环境最常用的方法。它的原理前面已经解释过通过设置环境变量在进程启动时就限定其能看到的GPU范围。这种方法的好处是彻底、干净代码本身无需任何修改只需要在启动命令前加上变量即可。如何使用在Linux/macOS的终端中# 只使用物理GPU 1 CUDA_VISIBLE_DEVICES1 python train.py # 使用物理GPU 0和2 CUDA_VISIBLE_DEVICES0,2 python train.py # 不使用任何GPU强制使用CPU CUDA_VISIBLE_DEVICES python train.py # 或者 CUDA_VISIBLE_DEVICES-1 python train.py在Windows的CMD中set CUDA_VISIBLE_DEVICES1 python train.py在Windows PowerShell中$env:CUDA_VISIBLE_DEVICES1; python train.py实操心得与陷阱脚本内获取在你的train.py脚本里torch.cuda.device_count()返回的已经是过滤后的数量索引从0开始。例如设置CUDA_VISIBLE_DEVICES2后脚本里torch.cuda.device(0)操作的其实就是物理GPU 2。子进程问题如果你在Python脚本中使用subprocess启动另一个需要GPU的进程子进程默认会继承父进程的环境变量。这通常是好事但如果你希望子进程使用不同的GPU就需要在subprocess.Popen中传入新的env字典来覆盖。与代码指定结合即使设置了环境变量代码里依然可以用torch.cuda.set_device(0)。此时这个0指的是虚拟索引。我个人的最佳实践是用环境变量做资源隔离用代码指定做逻辑控制。例如在K8s或Slurm集群中通过作业提交系统设置CUDA_VISIBLE_DEVICES在代码内部则基于这个虚拟视图来编写多卡逻辑。3.2 方法二使用torch.cuda.set_device单卡场景这个方法用于在Python代码内部为当前进程或更准确地说是当前CUDA上下文设置默认的GPU。它只影响设置之后创建的CUDA张量和模型。import torch # 将默认GPU设置为物理索引为1的卡 torch.cuda.set_device(1) # 现在所有.cuda()或.to(‘cuda’)操作都会默认使用GPU 1 x torch.tensor([1.0, 2.0]).cuda() # 张量x在GPU 1上 model MyModel().cuda() # 模型参数在GPU 1上 # 你可以通过以下方式验证 print(fCurrent device: {torch.cuda.current_device()}) # 输出: 1注意事项set_device的影响是线程局部的。如果你使用了多线程每个线程可能需要单独设置。它只设置“默认”设备。你仍然可以通过.to(‘cuda:0’)这样的方式明确指定其他可见的设备。但在单卡训练脚本中在开头写一句torch.cuda.set_device(device_id)是非常清晰的做法。如果CUDA_VISIBLE_DEVICES只指定了一张卡如2那么torch.cuda.set_device(0)就是操作那张卡。3.3 方法三使用.to(device)显式移动最灵活、最明确这是最精细、最推荐在模型和数据上使用的控制方法。它显式地将张量或模型移动到指定的设备上代码意图一目了然。import torch # 定义设备 device torch.device(‘cuda:1‘ if torch.cuda.is_available() else ‘cpu‘) # 或者更直接地 device_id 1 device torch.device(f‘cuda:{device_id}‘) # 将模型移动到指定设备 model MyModel().to(device) # 将数据移动到指定设备 data torch.randn(16, 3, 224, 224).to(device) labels torch.randint(0, 10, (16,)).to(device) # 在训练循环中确保每个batch的数据都到了正确的设备上 for inputs, targets in dataloader: inputs, targets inputs.to(device), targets.to(device) # ... 前向传播、计算损失、反向传播为什么这是最佳实践设备无关性通过将device定义为一个变量你可以轻松地在CPU和GPU之间切换或者通过改变device_id来切换GPU而无需修改模型和数据加载的逻辑。意图清晰任何阅读代码的人都能立刻知道每个张量应该在哪个设备上这对于调试内存溢出OOM或设备不匹配错误至关重要。支持多设备你可以为模型的不同部分指定不同的设备虽然不常见灵活性极高。3.4 方法四在DataLoader中使用pin_memory加速重要技巧这虽然不是直接“指定”GPU但却是GPU训练中与数据搬运密切相关的关键优化。当DataLoader的pin_memoryTrue时数据加载器会将数据张量固定在CPU的**页锁定内存Pinned Memory**中。原理与好处通常CPU内存中的数据需要先拷贝到一个临时缓冲区才能通过PCIe总线传输到GPU。而页锁定内存允许GPU的DMA引擎直接访问省去了中间拷贝的步骤从而实现从CPU到GPU的异步、高速传输。当你的数据预处理是瓶颈时这能显著提升训练速度尤其是当GPU比较“饿”等待CPU喂数据的时候。如何使用from torch.utils.data import DataLoader, Dataset train_dataset MyDataset(...) train_loader DataLoader( train_dataset, batch_size64, shuffleTrue, num_workers4, # 多进程加载数据 pin_memoryTrue, # 关键设置当使用GPU时开启 persistent_workersTrue # 如果num_workers0建议设为True以避免重复创建进程的开销 )重要警告pin_memoryTrue只在你使用GPU训练时才应该开启。它会增加CPU内存的占用。如果你在CPU上运行不仅没有收益反而会浪费内存。一个常见的模式是根据设备动态设置pin_memory (device.type ‘cuda‘) train_loader DataLoader(..., pin_memorypin_memory, ...)4. 多GPU训练DataParallel/DistributedDataParallel中的GPU指定当你需要将模型复制到多张GPU上并行训练以处理更大的批量或加速训练时指定GPU的逻辑需要升级。4.1 使用DataParallelDPDataParallel是单进程、多线程的并行方式使用非常简单但存在GPU负载不均衡和效率问题通常只用于快速原型。import torch.nn as nn # 假设我们想使用GPU 0, 1, 2 device_ids [0, 1, 2] # 1. 将模型放到第一张GPU上device_ids[0] model MyModel().to(f‘cuda:{device_ids[0]}‘) # 2. 用DataParallel包裹模型 model nn.DataParallel(model, device_idsdevice_ids) # 之后你只需要将输入数据放到第一张GPU上 # DataParallel会自动在forward过程中分割数据并分发到各GPU收集结果。 inputs inputs.to(f‘cuda:{device_ids[0]}‘) outputs model(inputs)DP的指定逻辑device_ids参数就是你告诉DataParallel要用哪些卡。它会以device_ids[0]作为“主设备”模型副本放在主设备上同时在其他设备上复制模型。每次前向传播时批次数据在主设备上被切分分发到各卡计算完成后再收集回主设备。因此你的数据只需要放到主设备上。4.2 使用DistributedDataParallelDDP工业标准DistributedDataParallel是多进程并行每个GPU对应一个独立的进程通信通过NCCL库效率远高于DP是当前多卡训练的事实标准。指定GPU的方式也完全不同。核心步骤使用torch.distributed.launch或torchrun启动脚本并为每个进程设置不同的LOCAL_RANK环境变量。在每个进程内部根据LOCAL_RANK来设置当前进程使用的GPU。示例脚本 (train_ddp.py)import torch import torch.distributed as dist import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP import argparse import os def main(): # 1. 解析参数获取本地进程排名 parser argparse.ArgumentParser() parser.add_argument(‘--local_rank‘, typeint, default-1) args parser.parse_args() # 2. 初始化进程组 dist.init_process_group(backend‘nccl‘, init_method‘env://‘) # ‘nccl‘是NVIDIA GPU间通信的最佳后端。 # 3. 根据local_rank设置当前进程使用的GPU # 这是最关键的一步每个进程独占一张GPU。 torch.cuda.set_device(args.local_rank) device torch.device(f‘cuda:{args.local_rank}‘) # 4. 创建模型并移动到当前设备 model MyModel().to(device) # 5. 用DDP包裹模型 # 注意device_ids参数通常只需指定为[local_rank]因为每个进程只有一张卡。 model DDP(model, device_ids[args.local_rank], output_deviceargs.local_rank) # 6. 准备数据每个进程会加载数据集的子集通过DistributedSampler train_sampler torch.utils.data.distributed.DistributedSampler(train_dataset) train_loader DataLoader(train_dataset, samplertrain_sampler, ...) # 7. 训练循环 for epoch in range(epochs): train_sampler.set_epoch(epoch) # 重要在每个epoch开始前打乱数据 for batch in train_loader: inputs, targets batch inputs, targets inputs.to(device), targets.to(device) # ... 训练步骤 # loss.backward()后DDP会自动同步所有进程的梯度。 if __name__ ‘__main__‘: main()启动命令# 使用 torch.distributed.launch (旧版仍可用) python -m torch.distributed.launch --nproc_per_node4 --use_env train_ddp.py # --nproc_per_node4 表示在本机的4张GPU上启动4个进程 # --use_env 会将 local_rank 通过环境变量传递脚本中需通过 os.environ[‘LOCAL_RANK‘] 获取 # 使用 torchrun (新版推荐更简洁) torchrun --nproc_per_node4 train_ddp.py # torchrun 会自动设置 local_rank 等环境变量DDP的GPU指定哲学在DDP中你不再需要手动指定一个device_ids列表。启动器launch或torchrun会负责生成多个进程并为每个进程分配一个唯一的local_rank通常是0, 1, 2, 3…。你的代码只需要根据传入的local_rank调用torch.cuda.set_device(local_rank)就能让每个进程绑定到对应的物理GPU上。这是一种“分而治之”的模式每个进程只关心自己那一片GPU天地。5. 模型推理部署时的GPU指定策略训练完成后模型上线提供推理服务。这时指定GPU的目标从“高效利用”变成了“稳定可控”。我们通常不希望一个推理服务占用所有GPU也不希望多个服务实例互相干扰。5.1 单模型单GPU服务这是最简单的场景。你可以在启动推理服务脚本时通过CUDA_VISIBLE_DEVICES将服务进程锁定在特定GPU上。# 启动一个Flask API服务只使用GPU 1 CUDA_VISIBLE_DEVICES1 python inference_api.py在inference_api.py内部代码可以写得很简单app Flask(__name__) device torch.device(‘cuda:0‘ if torch.cuda.is_available() else ‘cpu‘) # 注意这里索引是0因为环境变量只可见一张卡 model load_model(‘path/to/model.pth‘).to(device).eval() app.route(‘/predict‘, methods[‘POST‘]) def predict(): data request.get_json() tensor preprocess(data).to(device) with torch.no_grad(): output model(tensor) return postprocess(output)5.2 单模型多GPU负载均衡对于高并发场景单个GPU可能成为瓶颈。你可以启动多个相同的推理服务进程每个绑定到不同的GPU然后在前端用Nginx等负载均衡器分发请求。# 在screen或systemd中启动多个进程 CUDA_VISIBLE_DEVICES0 python inference_api.py --port 8000 CUDA_VISIBLE_DEVICES1 python inference_api.py --port 8001 CUDA_VISIBLE_DEVICES2 python inference_api.py --port 8002 5.3 多模型共享GPU服务器当一台服务器需要部署多个不同的模型时情况更复杂。你需要考虑显存隔离。简单的环境变量隔离可能不够因为一个模型即使空闲也会占用显存。更高级的做法是使用NVIDIA MPSMulti-Process Service或NVIDIA Triton Inference Server。NVIDIA MPS允许多个CUDA进程共享GPU上下文减少上下文切换开销并更精细地共享计算资源。但对于显存隔离它能力有限。NVIDIA Triton这是工业级的方案。Triton Server可以同时加载多个模型并允许你为每个模型实例指定在哪个GPU上运行、使用多少百分比的GPU计算资源。它还能动态批处理请求极大提升吞吐量。配置是通过一个config.pbtxt文件完成的例如instance_group [ { count: 1 # 实例数量 kind: KIND_GPU gpus: [ 0, 1 ] # 这个模型可以部署在GPU 0或1上 } ]对于大多数团队如果推理服务达到一定复杂度直接上Triton是比手动用Python脚本管理更明智的选择。6. 常见问题、调试技巧与性能考量6.1 典型错误与排查“RuntimeError: Expected all tensors to be on the same device”原因进行运算的张量不在同一个设备上比如一个在CPU一个在GPU 0或者一个在GPU 0一个在GPU 1。排查在错误发生前打印所有相关张量的.device属性。解决确保在数据送入模型前使用.to(device)将所有张量移动到统一设备。在复杂的自定义层或损失函数中要格外小心。“CUDA error: out of memory”原因显存不足。可能是批次太大、模型太大、或产生了中间变量缓存如梯度。排查在训练循环开始前和每个epoch结束后使用torch.cuda.memory_allocated(device_id)和torch.cuda.memory_reserved(device_id)监控显存。使用torch.cuda.empty_cache()可以释放PyTorch缓存的一些未使用的显存但这通常是治标不治本。解决减小batch_size。使用梯度累积gradient accumulation小批次计算梯度多次累积后再更新权重模拟大批次效果。使用混合精度训练AMP用torch.cuda.amp自动将部分计算转为float16减少显存占用和加速计算。检查是否有不必要的张量被长期引用如存储在列表里导致无法释放。“RuntimeError: CUDA error: invalid device ordinal”原因代码中指定的GPU索引超出了当前进程可见的范围。排查检查torch.cuda.device_count()确保你尝试访问的索引i满足0 i device_count。回想是否设置了CUDA_VISIBLE_DEVICES导致物理索引被重映射。多卡训练时只有一张卡在跑其他卡显存占用很低原因DP可能是数据批次batch太小。DP在主卡上切分批次如果总批次大小小于GPU数量有些卡就分不到数据。解决增大batch_size使其是GPU数量的整数倍。原因DDP可能DistributedSampler工作不正常或者每个进程的数据加载器配置有误。解决确保正确使用了DistributedSampler并且每个epoch开始时调用了sampler.set_epoch(epoch)。6.2 性能优化要点数据加载是瓶颈如果GPU利用率波动很大经常降到0%很可能是CPU数据预处理跟不上了。解决方案增加DataLoader的num_workers通常设为CPU核心数、使用pin_memoryTrue、将预处理尽可能移到GPU上如用torchvision.transforms的GPU加速操作或者使用更高效的图像解码库如turbojpeg。通信是瓶颈多卡训练DDP的梯度同步需要通信。如果模型很小而通信频繁通信开销可能占大头。解决方案对于超大模型考虑使用更高级的并行策略如模型并行、流水线并行。对于一般模型确保使用nccl后端并且服务器内GPU间使用NVLink或PCIe 4.0等高速互联。内核启动开销频繁启动大量微小的CUDA核函数会有开销。解决方案尽量使用PyTorch内置的、优化过的算子避免在循环中调用大量细粒度的自定义CUDA操作。6.3 监控工具推荐命令行nvidia-smi是最基本的。watch -n 1 nvidia-smi可以每秒刷新。可视化NVIDIA Nsight Systems: 提供系统级性能分析能看到CPU、GPU的时间线找出瓶颈所在。PyTorch ProfilerPyTorch内置的性能分析工具可以与TensorBoard集成可视化模型每个操作的耗时和显存消耗。with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3, repeat2), on_trace_readytorch.profiler.tensorboard_trace_handler(‘./log‘), record_shapesTrue ) as prof: for step, data in enumerate(train_loader): if step (1 1 3) * 2: # 对应schedule的循环 break train_step(data) prof.step()运行后使用tensorboard --logdir./log查看分析结果。指定GPU这个操作从简单的环境变量到复杂的分布式训练集成贯穿了深度学习工程化的始终。我的体会是在实验阶段就养成好习惯总是显式地使用.to(device)并通过启动脚本来控制GPU资源。这样当你的代码从笔记本迁移到服务器从单卡扩展到多卡时会减少大量的适配和调试时间。记住清晰的设备管理策略是构建稳定、可复现、可扩展的深度学习项目的基石。最后一个小技巧在你项目的README.md里明确写下运行所需的环境变量和启动命令这对你的合作者和未来的自己都是极大的帮助。