PyTorch GPU 训练实战:从零搞懂模型、数据如何统一到同一设备 PyTorch GPU 训练实战从零搞懂模型、数据如何统一到同一设备副标题手把手带你排掉 GPU 训练的五大坑——设备不一致、忘记迁移、跨设备加载、显存爆炸、多卡并行怎么选你兴致冲冲地斥巨资换了张显卡想体验一把十倍加速的快感。结果一行model(x)还没跑完PyTorch 就甩给你一句冰冷的报错RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!模型在 GPU 上数据却还在 CPU 上——它们俩隔了一道内存墙根本没法直接做运算。这几乎是所有新手上 GPU 的成人礼。前面第 43 篇我们讲过Tensor 有个device属性它决定这个张量在哪台设备上参与计算。今天我们就把整个训练流程从 CPU 搬到 GPU把模型、数据、损失函数怎么统一到同一设备这件事讲得明明白白顺便把多卡并行、混合精度、显存爆炸这些进阶坑也一并扫清。我会用「仓库—车间」的比喻串起全文把 CPU 想成大仓库容量大但搬运慢把 GPU 想成高速车间算力猛但工位有限。训练时你当然希望原料数据和机器模型都在同一个车间里否则就得不停地跨厂房搬货既慢又容易出错。一、先确认你的车间到底开没开搬设备之前先别急着写.to()。第一步永远是确认机器上到底有没有可用的 GPU——万一你在一台纯 CPU 的服务器上写了.to(cuda)代码会直接崩。importtorchprint(torch.cuda.is_available())# True 表示装了 CUDA 版 PyTorch 且有可用 N 卡print(torch.cuda.device_count())# 显卡数量比如你有 4 张卡这里就是 4print(torch.cuda.get_device_name(0))# 第 0 张卡的名字比如 NVIDIA RTX 4090但真实项目里你写的代码大概率要在有无 GPU 的机器上都跑得通比如本地没卡、服务器有卡。所以别把cuda写死而是用一个万能写法自动选设备devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)print(f使用设备:{device})# 有卡 → cuda没卡 → cpu直觉补充torch.device只是个地址标签本身不占资源。真正吃显存的是你把张量/模型.to(device)搬过去的那一刻。所以先定义device再在必要处搬运是 GPU 代码的黄金起手式。 点击在线体验此交互组件二、GPU 训练的铁律三样东西必须同处一室GPU 训练只有一条核心原则记住这句话就够了模型、输入数据、标签必须在同一个设备上。用.to(device)把它们搬过去缺一不可。为什么是三样而不是两样我们拆开看。1. 模型把机器搬进车间model.to(device)会把模型里所有可学习参数weight、bias都迁移到目标设备。注意它是一个就地in-place且有返回的操作两种写法等价modelMyModel()modelmodel.to(device)# 推荐显式接住返回值语义更清晰# 或者model.to(device)# 也行参数已被原地迁移2. 输入数据每一批原料都要搬模型搬过去了原料还在仓库里可不行。训练循环里每一个 batch 的x和y都要单独搬forx,yintrain_loader:xx.to(device)# ⚠️ 输入数据迁移yy.to(device)# ⚠️ 标签也要迁移漏了这一个照样报错...⚠️ 新手最常犯的错就是模型搬了数据忘了搬。记住——DataLoader产出的数据默认永远在 CPU模型搬去 GPU 并不会顺手把数据也搬过去。3. 损失函数多数情况自动跟随CrossEntropyLoss、MSELoss这些标准损失函数没有可学习参数它们的运算发生在模型输出和标签所在的设备上所以只要logits和y都在 GPU损失函数自然就在 GPU 上算不需要你手动搬。但有个例外如果你自定义了带参数的 loss比如内部含nn.Parameter那这个 loss 模块也得.to(device)否则它的参数在 CPU、输入在 GPU照样炸。 点击在线体验此交互组件三、完整 GPU 训练模板直接抄把上面的原则拼起来就是一份能直接上手的 GPU 训练骨架。我特意把关键迁移点用 ⚠️ 标出importtorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoader# 1. 设备有无 GPU 自动兼容devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)# 2. 数据pin_memoryTrue 让 CPU→GPU 拷贝更快后文详解train_datasetMyDataset(...)train_loaderDataLoader(train_dataset,batch_size64,shuffleTrue,num_workers4,pin_memoryTrue)# 3. 模型迁移到设备modelMyModel().to(device)# 4. 损失和优化器criterionnn.CrossEntropyLoss()optimizertorch.optim.Adam(model.parameters(),lr0.001)# 5. 训练循环forepochinrange(epochs):model.train()forx,yintrain_loader:xx.to(device)# ⚠️ 数据迁移yy.to(device)# ⚠️ 标签迁移optimizer.zero_grad()logitsmodel(x)# 模型和数据都在 GPU运算丝滑losscriterion(logits,y)loss.backward()optimizer.step()️易混淆辨析model.to(device)和x.to(device)不是一回事。前者搬的是模型参数持久状态“只需在训练前做一次后者搬的是每个 batch 的临时数据瞬时状态”要在循环里每个 batch 都做。很多人以为模型搬了数据自动跟着搬——并不会。四、三大高频报错与对症下药报错 1模型和数据不在同一设备最经典的RuntimeError: Expected all tensors to be on the same device。排查只需两行打印print(next(model.parameters()).device)# 模型参数在哪 → 比如 cuda:0print(x.device)# 数据在哪 → 比如 cpu两者必须完全一致。修复就是确保都.to(device)。报错 2评估时忘记迁移数据这个坑特别阴——训练时你小心翼翼地写了x.to(device)一到测试阶段复制粘贴代码时把那两行漏了# ❌ 训练时迁移了评估时忘了model.eval()withtorch.no_grad():forx,yintest_loader:# x 还在 CPUlogitsmodel(x)# 报错模型在 GPU数据在 CPU修复评估循环里也要把数据搬过去和训练一模一样model.eval()withtorch.no_grad():forx,yintest_loader:xx.to(device)yy.to(device)logitsmodel(x)报错 3保存的模型加载到无 GPU 的机器你在有卡的机器上保存同事在无卡的机器上加载直接崩# 保存时在 GPUtorch.save(model.state_dict(),model.pth)# 加载到 CPU 的机器上model.load_state_dict(torch.load(model.pth))# 报错参数在 cuda:0当前没 GPU修复加载时指定map_location强行把参数映射到目标设备# 想加载到 CPUmodel.load_state_dict(torch.load(model.pth,map_locationcpu))# 想加载到当前可用设备有卡上卡无卡上 CPUmodel.load_state_dict(torch.load(model.pth,map_locationdevice))state_dict本质是个普通的 Python 字典键是参数名值是张量它不绑定设备信息。map_location只是告诉 PyTorch把这些张量搬到哪并不改变你保存时的逻辑。这是个基础但极容易忘的点。五、让 GPU 真正吃饱四个加速技巧把模型和数据搬上 GPU 只是第一步。很多人搬上去后发现显存占了速度却没快多少——因为 GPU 在空等数据。下面四招让它真正转起来。1.pin_memoryTrue给数据走高速通道DataLoader(dataset,pin_memoryTrue)普通内存pageable memory在交给 GPU 前CUDA 还得先把它拷到一块锁页内存pinned memory再传。设了pin_memoryTrue数据加载时就直接放在锁页内存里省掉这一次额外拷贝CPU→GPU 传输更快。GPU 训练基本是必开项。2.num_workers 0多进程预加载别让 GPU 闲着DataLoader(dataset,num_workers4)数据预处理解码、增强、归一化在 CPU 上做。开多个子进程并行准备下一个 batchGPU 算完当前 batch 时下一批已经候着了。上一篇我们详细讲过它的原理这里记住结论GPU 利用率低先加num_workers。3. 用nvidia-smi看 GPU 到底忙不忙nvidia-smi如果 GPU 利用率长期低于 30%说明瓶颈在数据加载GPU 大部分时间在发呆。优先加num_workers、开pin_memory、简化数据增强。4. 混合精度训练进阶必杀前面的代码全程用float32。混合精度AMP让前向和反向的大部分计算用float16只在容易溢出的地方保留float32fromtorch.cuda.ampimportautocast,GradScaler scalerGradScaler()# 防止 float16 梯度下溢的缩放器forx,yintrain_loader:x,yx.to(device),y.to(device)optimizer.zero_grad()withautocast():# 自动把能转 float16 的计算转掉logitsmodel(x)losscriterion(logits,y)scaler.scale(loss).backward()# 先放大损失避免梯度下溢scaler.step(optimizer)# 再缩回来更新scaler.update()显存大约减半速度提升约 1.5–2 倍是大模型训练标配。入门阶段可以先不碰但要知道它有这回事。六、显存爆炸OOM与排查当你把batch_size调大想喂饱 GPU很可能迎来第二个经典报错RuntimeError: CUDA out of memory. Tried to allocate ...常见原因batch_size太大、模型太大、评估时忘了torch.no_grad()导致计算图堆积显存泄漏、或者多卡时每卡都拷了一份完整模型。修复清单减batch_size如 64 → 32——最直接有效。评估/推理时务必加with torch.no_grad()不保留计算图。调用torch.cuda.empty_cache()清理显存碎片治标不治本但有时能救急。模型太大就换小模型或上混合精度显存直接减半。# 评估时正确姿势no_grad 是关键model.eval()withtorch.no_grad():# ← 忘这个显存会一路涨forx,yintest_loader:x,yx.to(device),y.to(device)logitsmodel(x)七、进阶多卡并行到底是怎么回事当你有一张卡还不够、或者手握多卡服务器时就到了并行训练。这里必须主动辨析一个常见误解❌ 误解“用多张 GPU模型就变得更聪明了。”✅ 真相多卡只改变算得有多快不改变学得好不好。它本质上是通过并行增大等效batch_size、缩短每个 epoch 的时间让你能更快试更多超参、喂更多数据——这才间接提升了最终效果。PyTorch 给你两条主流路径方案模式原理适用DataParallel(DP)单进程、多卡主卡把一份 batch拆成多份分给各卡各卡前向反向后梯度回传主卡求平均再更新快速上手、2~4 卡、卡间通信快同一台机DistributedDataParallel(DDP)多进程、多卡每张卡一个独立进程各自持完整模型副本反向后通过 Ring AllReduce 同步梯度无主卡瓶颈多机多卡、大规模训练、性能首选DataParallel的写法最简单一行包裹即可# 多卡时把模型包一层单进程主卡 cuda:0 做梯度汇总iftorch.cuda.device_count()1:modelnn.DataParallel(model)modelmodel.to(device)但 DP 有个致命性能短板所有梯度都要汇总到主卡cuda:0做平均主卡容易成为瓶颈且只支持单机。现代大规模训练几乎都推荐DDP虽然写法稍复杂要初始化进程组、设置local_rank但每张卡各自算、各自同步没有单点瓶颈。️辨析DP 是参数服务器思路主卡分发汇总DDP 是对等同步思路Ring AllReduce。新手先会用 DP 理解数据并行的概念即可真要上规模请直接学 DDP。 点击在线体验此交互组件八、三个高频错误收藏备用错误原因解法Expected all tensors to be on the same device模型在 GPU、数据在 CPUmodel.to()与x.to()、y.to()成对出现加载模型报错参数在 cuda:0保存于有卡机、加载于无卡机torch.load(path, map_locationcpu)CUDA out of memorybatch 太大/忘 no_grad/显存泄漏减 batch、加no_grad、清缓存、混合精度错误 1只迁移模型不迁移数据modelmodel.to(device)# 忘了 x.to(device)logitsmodel(x)# 报错数据还在 CPU错误 2在 GPU 上创建 Tensor 时用了 CPU 默认# ❌ 数据在 CPUxtorch.randn(32,10)modelmodel.to(cuda)model(x)# 报错# ✅ 创建时就指定设备省一次拷贝xtorch.randn(32,10,devicecuda)错误 3OOM见第六节核心是减batch_size 评估加torch.no_grad() 必要时empty_cache()。核心要点小结GPU 训练铁律模型、输入、标签必须同处 CPU 或同处 GPUmodel.to(device)和x.to(device)/y.to(device)要成对写。设备探测先行用torch.device(cuda if torch.cuda.is_available() else cpu)兼容有无 GPU标准损失函数无需手动搬。评估别漏迁移测试循环里同样要把数据搬上设备并加torch.no_grad()。跨设备加载用map_locationtorch.load(path, map_locationdevice)可把模型参数映射到任意设备。喂饱 GPUpin_memoryTruenum_workers0消除数据瓶颈混合精度AMP显存减半、提速近 2 倍。OOM 先减 batch评估忘no_grad是显存泄漏元凶。多卡只加速不增智DP 有主卡瓶颈、适合快速上手DDP 对等同步、适合规模训练——两者都是数据并行不改变模型上限。动手思考题写一个小函数train_one_epoch(model, loader, device, optimizer, criterion)在 GPU 上训练一个 epoch要求同时正确处理model.train()、数据迁移和梯度清零。下面代码哪里有问题运行时大概率报什么错devicecudamodelmodel.to(device)forx,yinloader:logitsmodel(x)losscriterion(logits,y)训练时遇到CUDA out of memory请列出至少三个修复方案并说明哪个最治本。DP 和 DDP 的本质区别是什么为什么说多卡不改变模型上限只改变迭代速度在评论区用你自己的话讲一遍我们一起纠正 下一篇我们聊正则化——Dropout、Early Stopping、Weight Decay怎么让模型不死记硬背训练集上线后照样稳。 关于本系列本文是「AI 学习路线 · 阶段四:PyTorch 深度学习基础」系列中的一篇。所有文章在我的个人博客上都有可交互动画 完整学习路线版本,建议配合食用 在博客上阅读本文原版(含可交互组件、公式动画) PyTorch GPU 训练实战从零搞懂模型、数据如何统一到同一设备️查看完整 AI 学习路线(从 0 到进阶,持续更新) bestsdz.xyz觉得有帮助的话,欢迎去博客点个收藏 ⭐,你的支持是我更新的最大动力!