先说个很多Pytorch新手都会遇到的场景本地CPU上跑得好好的训练代码一放到带GPU的机器上就开始连环报错。一翻日志要么是Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!要么是莫名其妙的显存暴涨。其实这些问题的根源通常只有一个模型和数据没有待在同一个设备上。Pytorch里负责这件事的就是那一串API——model.to(device)、model.cuda()、model.cpu()、tensor.to(device)再往多卡走就是DataParallel。这篇文章把这些API的底层逻辑、正确用法和常见坑完整过一遍适合刚接触Pytorch的新手也适合单卡跑得挺顺、但一看多卡代码就发懵的同学。我会尽量把“为什么要这么写”讲透而不只是给一段能跑的代码。1. 设备和张量的底层逻辑为什么总会碰到device不匹配1.1 CPU、GPU与内存、显存的关系先明确一个底层事实Pytorch中的张量Tensor是有“居住地”的。默认情况下你用torch.tensor()或者torch.randn()创建出来的张量放在CPU内存里而GPU上能用的空间叫显存两者是物理隔离的。CPU内存由操作系统管理显存由CUDA驱动管理二者之间通过PCIe总线进行拷贝速度远低于各自内部的读写速度。这意味着什么意味着你在CPU上创建了一个张量然后直接拿它和一个在GPU显存里的张量做加法Pytorch会直接报错。这不是设计缺陷而是底层硬件决定的——CPU无法直接访问显存地址GPU也无法直接访问内存地址在不支持统一内存的架构上。所以任何跨设备的运算都必须先把数据显式地拷贝到同一个设备上。模型也是一样的。一个nn.Module里包含的参数nn.Parameter和缓冲区buffer比如BatchNorm里的running_mean、running_var本质上都是张量。当你把一个模型从CPU搬到GPU时真正做的事就是把这些参数和缓冲区的存储位置整体迁移过去。1.2 为什么device不匹配是新手最常见的问题打个比方CPU内存是办公室的桌面GPU显存是一个只能在机房才能访问的保险柜。你可以把文件放在桌面上也可以把文件锁进保险柜但你不能坐在工位上直接修改保险柜里的文件也不能在机房外面凭空读取桌面上的文件。Pytorch里的model.to(device)和tensor.to(device)干的就是“把文件从桌面搬进保险柜”或者“从保险柜搬回桌面”这种事。很多新手犯的错误是在一台只有CPU的电脑上写代码时所有数据都自然在CPU上模型也在CPU上一切正常。到了有GPU的机器上把device设置成了cuda模型也调用了model.cuda()但忘了把每个batch的输入数据和标签也搬过去。于是模型在显卡上等着数据还在内存里两边一对上就报错。理解了张量和参数都有device属性之后这类问题就比较好排查了。判断当前环境有没有可用的GPU标准写法是import torch device torch.device(cuda if torch.cuda.is_available() else cpu) print(device) # 有GPU时输出 cuda否则输出 cputorch.device对象是Pytorch推荐的设备表示方式它同时在内部携带了typecpu/cuda和indexGPU编号。torch.cuda.is_available()用来探测CUDA环境是否可用这在写跨平台训练脚本时是必须的一步。1.3 tensor的dtype和device是两套独立属性这里要特别区分两个概念dtype数据类型和device设备类型。一个张量可以是float32、float64、int64等同时它也可以位于cpu或cuda上。两者相互独立但都会影响运算是否正确。常见的坑之一是在某些CPU环境下你可能会为了提速把数据转成float16但GPU上用float16配合混合精度训练时需要额外处理梯度缩放否则模型可能不收敛。这是后面章节会展开的内容这里先记住一点先把dtype理清楚再把device理清楚一个一个来不要混在一起调。2. model.to(device) 与 model.cuda() / model.cpu() 的完整对比2.1 三个API的基本用法Pytorch中迁移模型有三个常见写法表面上看功能重叠但使用场景和推荐程度有差别。API作用对象等价关系返回值常用场景model.to(device)nn.Module通用写法可指定设备self原对象既支持CPU也支持GPU推荐使用model.cuda()nn.Module等价于model.to(torch.device(cuda))self仅用于GPU早期代码常见model.cpu()nn.Module等价于model.to(torch.device(cpu))self把模型从显存放回内存model.to(device)里的device可以是一个torch.device对象也可以是一个设备字符串比如cpu、cuda、cuda:0。这里有个容易忽略的细节model.to(cuda)会默认使用编号为0的GPU如果想用其他卡需要写成model.to(cuda:1)。我建议的处理方式是只维护一个device变量全代码统一用model.to(device)不要把cuda()和cpu()散落在各处。为什么因为代码可移植性。如果某天你需要在没有GPU的机器上做调试只要把最上面一行device的定义改掉整份代码的模型迁移都会跟着变。但如果你到处写死了model.cuda()在纯CPU环境下就会直接报错得全局搜索替换非常痛苦。2.2 迁移的底层操作参数和缓冲区被就地替换很多教程会告诉你“把模型搬到GPU”但它没说清楚“搬”这个动作到底做了什么。实际上nn.Module.to()以及等价的cuda()和cpu()会遍历模型内部的所有参数和缓冲区逐个调用tensor.to(device)然后把这些新设备上的张量重新赋回模型的参数字典里。因为整个操作是就地修改的所以你不需要写model model.to(device)直接调用model.to(device)就够了。这里有一个非常容易踩的经典坑模型迁移和优化器创建的先后顺序。看看下面这段错误示例model MyModel() optimizer torch.optim.Adam(model.parameters(), lr0.001) model model.cuda() # 或者 model.to(cuda)问题出在哪optimizer在创建时会把model.parameters()返回的参数对象存到自己的param_groups里。如果你之后再执行model.cuda()模型内部的参数可能被替换成新的GPU张量对象但optimizer里保存的依然是迁移前的那批CPU参数对象。结果就是模型在GPU上计算并产生梯度但优化器更新的却是CPU上的旧参数或者直接报device不匹配。推荐做法是先迁移模型再创建优化器model MyModel().to(device) optimizer torch.optim.Adam(model.parameters(), lr0.001)这个顺序问题在单卡上偶尔能跑通但一旦涉及多卡或混合精度各种诡异问题都会冒出来。建议从一开始就养成固定顺序的习惯先定device再建模型并迁移最后创建优化器。2.3 别忘了bufferBatchNorm统计量的迁移model.to(device)对参数的迁移大家都熟悉但很多人会忘记模型里的buffer也是需要迁移的。最常见的buffer就是BatchNorm层的running_mean和running_var它们记录了整个训练过程中每个通道的均值和方差不是通过反向传播更新的而是通过滑动平均维护的。好消息是nn.Module.to()默认也会把buffer一起迁移所以常规的model.to(device)不会漏掉它们。但如果有人手动创建了不含在nn.Module内部的张量、并且单独引用此时就需要手动.to(device)。例如在模型外维护一个EMA均值或队列的时候必须记得对它们同步做设备迁移。否则你可能会遇到“模型在GPU上buffer在CPU上”这种贴别明显的报错而且很难通过看模型结构发现。3. 数据的设备迁移训练循环里的正确姿势3.1 每个batch都要遵守的“搬运三步”模型迁移完成只是第一步数据侧同样需要迁移。一个标准的训练循环长这样for epoch in range(num_epochs): for batch_idx, (images, labels) in enumerate(train_loader): images images.to(device, dtypetorch.float32) labels labels.to(device, dtypetorch.long) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()看到没images和labels都必须显式调用.to(device)。这里很多人会问为什么Pytorch不自动把数据搬过去因为Pytorch的设计哲学是“显式优于隐式”。如果你希望数据在CPU上做预处理然后只在送入模型前搬到GPU这种灵活性只有通过手动控制才能实现。自动迁移反而会让数据流不可控。另外注意dtype的匹配。图像数据进入卷积网络通常需要float32而标签进入损失函数通常是int64CrossEntropyLoss要求target必须是long类型。在部分旧版本或半精度训练场景下dtype不匹配会直接报错所以我在代码里习惯把dtype和device写在同一个.to()调用里一步到位。3.2 pin_memoryTrue到底提升了什么DataLoader里有一个参数叫pin_memory很多教程说“设成True能加速”但没说清楚原理。默认情况下CPU内存中的张量是分页内存pageable memory操作系统可以随时调整这些内存页的位置。而GPU需要拷贝CPU数据时走的是DMA直接内存访问它需要和一片固定的物理内存打交道。如果数据在分页内存里CUDA会先把数据拷贝到一个临时锁页缓冲区然后再通过DMA传到显存多了一步。pin_memoryTrue的意思是在加载数据时直接把数据放进锁页内存page-locked memory这样向GPU拷贝时可以直接走DMA无需中间缓冲。这在高吞吐训练场景下确实能减少一部分host到device的拷贝时间。train_loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue)但要注意pin_memory只是为.to(device)拷贝提速并不会自动把数据放到GPU上数据仍然在CPU锁页内存里。而且num_workers0时每个worker子进程复制数据也需要考虑锁页内存的开销。通常情况下pin_memoryTrue是值得开的但如果你发现数据加载内存占用异常高也可以关掉对比一下。3.3 在loss累计和显存释放上的常见失误还有一个经常被忽略的显存问题跟设备迁移无关但彼此纠缠。很多人写训练循环时习惯用total_loss loss来累计loss这是错的因为loss本身是GPU上的张量每次累加都会生成一个新的GPU张量导致显存持续增长直到出现OOM。正确做法是total_loss loss.item()。.item()会把GPU上的标量张量转换为Python的float数然后在CPU上累计完全不占显存。这个细节和设备迁移是两件事但在排查显存问题时经常一起出现。建议在训练循环里统一遵守两条规则需要反向传播的loss保持为GPU张量需要记录或打印的loss用.item()转出来。3.4 标签到底要不要迁移不少人认为标签不用搬到GPU理由是损失函数会自动处理。实际上CrossEntropyLoss在计算时要求input和target必须在同一个设备上。如果input是GPU张量target是CPU张量执行criterion(outputs, labels)时会直接报device不匹配。所以标签也必须跟着迁移。少数情况下标签的迁移可以在损失函数内部规避比如自己写损失时用labels.to(outputs.device)临时搬一下。但我还是建议在循环开头统一迁移干净这样整个训练过程的设备状态一目了然不用在后面的代码里到处去找设备来源。4. DataParallel多卡并行从单卡到多卡的提速与踩坑4.1 DataParallel的工作原理单卡训练顺畅之后很多人会想用多张GPU提速。nn.DataParallel是最容易上手的方案它做的事情可以概括为在每次前向传播时把输入数据在batch维度上切成多份分发给多个GPU上的模型副本并行计算然后把所有GPU的输出重新拼接起来在主卡上统一计算损失和梯度。数据流大概是这样的主卡device_ids[0]接收完整batch的输入数据。把输入在batch维度切分分发给所有GPU。每张GPU上有一个模型副本各自对这个子batch执行前向计算。所有GPU的输出收集到主卡按batch维度拼接。主卡计算损失反向传播时梯度自动聚合到主卡并更新主卡模型参数。下一轮参数同步到其他GPU。注意一点Pytorch官方后来更推荐用DistributedDataParallelDDP做多卡训练因为DataParallel在实现上受制于Python线程和GIL而且主卡经常成为瓶颈。但如果只是简单实验、不想引入复杂的启动脚本DataParallel依然够用前提是你知道它的坑在哪。4.2 DataParallel的标准写法device torch.device(cuda if torch.cuda.is_available() else cpu) model MyModel().to(device) if torch.cuda.device_count() 1: model nn.DataParallel(model)这里有一个很关键的细节DataParallel包装之后模型对象变成了nn.DataParallel实例它内部持有一个module属性指向原始的MyModel。所以如果你需要在代码里访问模型的原始结构比如打印层、提取特征图要用model.module。另外在使用DataParallel时输入数据仍然要先搬到主卡所在的设备images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels)device取的是当前默认CUDA设备通常cuda:0。除非你明确指定了device_ids否则不要试图把模型放在cuda:1然后指望DataParallel自动处理很多诡异的device不匹配都是这么来的。4.3 损失函数的位置最容易出错的点用DataParallel之后损失函数的写法直接决定你的loss值是否正确。这里有个很多人踩过的坑把损失计算写在了模型的forward()里并返回一个标量loss。DataParallel有个特性如果模型返回的是一个标量张量比如loss它会这个标量的聚合方式处理具体行为在不同版本中曾有差异。结果是你在外部拿到的loss可能是各GPU上loss的某种聚合结果而不是完整batch的平均loss数值和单卡对不上。更稳妥的做法是让模型的forward只返回输出损失在外部计算。这样DataParallel把所有GPU的输出拼成一个完整batch的output再交给损失函数得到的就是整个batch正确的平均损失。outputs model(images) loss criterion(outputs, labels)如果模型是DataParallel包装的outputs已经按batch维度拼接好criterion会正常接受它。这段代码无论单卡还是多卡都能跑而且数值一致。我强烈建议把“前向输出”和“损失计算”分开设计不仅是为了DataParallel也便于后续切换DDP或混合精度。4.4 主卡显存倾斜和BatchNorm问题DataParallel最大的物理瓶颈是主卡显存会明显偏高。原因是主卡不仅要存一份完整模型参数还要承担输入数据的接收、切分、输出聚合等工作。如果你的batch size比较大主卡显存可能先爆掉其他卡却还有余量。常见的缓解手段是减小batch size、让批大小能被卡数整除、尽量不要在主卡上做额外的显存消耗操作。但说白了DataParallel的扩展性有限这也是为什么它不适合大规模训练。如果你的需求是4卡以上、或者每张卡显存都很紧张建议直接换DistributedDataParallel。还有一个容易被忽略的点是BatchNorm在DataParallel下的行为。DataParallel把batch切分到多张卡后每张卡上的BatchNorm是在自己的子batch上计算均值和方差的相当于有效batch size变成了原来的1/N。原本单卡batch644卡并行后每张卡实际只看到16个样本对BatchNorm统计量的估计会变差甚至影响收敛。如果模型包含BatchNorm且batch size本身不大可以考虑切换到SyncBatchNorm它会让所有卡同步计算全局的均值和方差model torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)不过SyncBatchNorm在多卡训练时需要确保能同步这通常配合DDP使用效果更好。在DataParallel场景下如果模型较大且本身就是大批量训练普通BatchNorm问题不大。4.5 DataParallel和保存加载的“module前缀”问题使用DataParallel后模型的state_dict的键名会多出一层module.前缀。这是最经典的多卡切换坑之一。比如原模型有一个conv1.weight键用DataParallel包装后保存出来的键名是module.conv1.weight。如果你把这份权重直接加载到一个没有经过DataParallel包装的模型里就会报size mismatch或Missing key(s)。应对方案有两种。第一种保存时保存原始模型的state_dicttorch.save(model.module.state_dict(), best_model.pt)这样保存出来的键名没有module.前缀单卡加载很直接。第二种如果你想统一保存包装后的模型可以在加载时手动去掉前缀from collections import OrderedDict checkpoint torch.load(best_model.pt, map_locationcpu) state_dict checkpoint[model_state_dict] new_state_dict OrderedDict() for k, v in state_dict.items(): name k[7:] if k.startswith(module.) else k # 去掉module.前缀 new_state_dict[name] v model.load_state_dict(new_state_dict)这个技巧我相信写过分布式训练脚本的人都有自己的一份因为只要在多卡和单卡之间倒腾过一次模型就一定会碰到。5. 模型保存与加载时的设备地图map_location和优化器状态5.1 map_location决定了checkpoint恢复到哪个设备torch.save保存模型时默认会把张量当前的设备信息一起存下来。如果你在GPU上保存了一个模型然后在一台没有GPU的机器上直接torch.loadPytorch会尝试把张量加载到cuda设备然后报错“Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False”。解决办法是使用map_location参数checkpoint torch.load(best_model.pt, map_locationcpu)这句代码相当于告诉Pytorch不管这份checkpoint当初在哪个设备上都把张量加载到CPU内存。这是最保险的做法。加载到CPU之后再通过model.load_state_dict把参数填进模型最后整体调用model.to(device)把模型迁移到需要的设备上。我看到有些代码在加载时直接map_locationcuda:0这当然也行但它把“加载”和设备绑定死了复现代码的人一旦GPU环境和你不一致又要改代码。所以我个人的习惯是永远先加载到CPU再整体迁移。5.2 优化器状态里的张量也要处理设备模型权重只是checkpoint的一部分如果你保存了optimizer.state_dict()里面还包含每个参数对应的动量、方差等状态这些同样是张量同样带着设备信息。完整的保存逻辑一般是这样checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, best_acc: best_acc, } torch.save(checkpoint, checkpoint.pt)加载的时候如果先map_locationcpu加载了checkpoint那么optimizer.state_dict()里的张量也都在CPU上。此时如果直接把optimizer.load_state_dict填给一个已经迁移到GPU的优化器对象可能会出现设备不匹配。稳妥的做法是先把优化器状态加载到CPU然后手动把内部张量迁移到目标设备。checkpoint torch.load(checkpoint.pt, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) # 将优化器内部的张量迁移到device for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] v.to(device)这段代码实际排查时很有用。不处理优化器状态的话会出现一种很迷惑的现象模型参数在GPU上正常运行但优化器step时报device mismatch或者训练了一段时间后loss异常。5.3 多卡checkpoint和单卡模型的兼容性处理如果你保存的是DataParallel包装后的模型同时你又希望这份权重能在单卡模型上加载那第4章说的module.前缀移除逻辑就派上用场。反之如果你保存的是单卡模型的state_dict希望加载到DataParallel模型里那需要手动加上module.前缀。写一个通用的加载函数可以避免每次都复制粘贴这些逻辑def load_state_dict_flexible(model, state_dict): prefix module. if all(k.startswith(prefix) for k in state_dict.keys()): state_dict {k[len(prefix):]: v for k, v in state_dict.items()} elif not all(hasattr(m, module) for m in model.modules()) and \ not any(k.startswith(prefix) for k in state_dict.keys()): # 模型是DataParallel但权重没有前缀需要补上 state_dict {prefix k: v for k, v in state_dict.items()} model.load_state_dict(state_dict)这段逻辑并不复杂核心就是检查state_dict里的键名是否以module.开头然后做相应的增删。有了这个函数无论你的checkpoint来自单卡还是多卡都能顺利加载。6. 常见错误排查与设备调试技巧6.1 高频报错速查表下面这个表格里列出的几类问题基本覆盖了90%以上和device相关的报错。我按错误信息的关键字整理了原因和解决办法。报错信息关键片段最常见原因解决办法Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu输入数据或标签没有迁移到GPU在训练循环里对images和labels调用.to(device)Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False在没有GPU的环境下直接加载GPU模型torch.load时加map_locationcpusize mismatch for ... conv1.weight模型结构与checkpoint不匹配通常是module.前缀检查并处理module.前缀CUDA out of memorybatch过大、主卡显存不足、loss累计张量等减小batch、使用.item()、检查显存分配device-side assert triggeredGPU上张量数值异常常见于标签越界或NaN检查标签取值范围、是否混用了dtypemodule must have its parameters and buffers on device cuda:0模型和输入设备不一致多卡时主卡选择错误统一主卡设备输入先to(device)这些报错本质上都在说同一件事某个张量或参数不在你期待的设备上。排查时先问自己三个问题模型在哪个设备输入在哪个设备标签在哪个设备三个设备都对上大多数问题就消失了。6.2 定位device问题时的实用工具如果你遇到一个奇怪的device相关报错但怎么都看不出来是哪个张量导致的有两个环境变量值得一试。第一个是CUDA_LAUNCH_BLOCKING1它强制让每个CUDA操作同步执行这样报错时可以给出更精确的调用栈而不是等到异步执行完才弹出一个模糊错误。第二个是TORCH_SHOW_CPP_STACKTRACES1它会显示C层面的堆栈信息对定位底层操作很有帮助。CUDA_LAUNCH_BLOCKING1 TORCH_SHOW_CPP_STACKTRACES1 python train.py这两个变量只在调试时开正式训练时不要开。因为它们会让GPU的异步执行变成同步执行训练速度会明显下降。另外如果你怀疑某个张量在哪个设备上可以直接打印print(model.device) # nn.Module没有device属性这会报错 # 正确做法是检查第一个参数 print(next(model.parameters()).device) print(images.device, labels.device)注意nn.Module本身没有.device属性网上偶尔会看到model.device的写法那是某些封装的库才有的。标准Pytorch里想获取模型的设备位置要么用next(model.parameters()).device要么用model.state_dict()里的某个张量。这个细节虽然小但排查问题时能省不少时间。6.3 我的调试习惯和一些经验最后分享几个我在实际调试中总结的习惯。第一所有训练脚本开头都统一放一个device定义后续所有模型和数据迁移都引用这个变量不写死cuda或cuda:0。这样在CPU上调试逻辑、在GPU上训练模型只需要改一行代码。第二搭建模型后立刻做一次简单的设备自检用一个小张量把模型跑一遍确认模型和数据在同设备上能正常前向再做正式训练。我见过太多人是看到报错才回头检查设备不如省下这一步。第三多卡训练时先用batch_size2的极小配置跑一个batch确认DataParallel的前向、损失、反向都能走通再放大batch_size。这么做可以快速区分是“代码逻辑问题”还是“显存容量问题”。多卡场景下的报错往往包含大量CUDA内部信息如果一上来就用大batch你可能分不清到底是显存不足还是代码里的device没配对。第四保存checkpoint时固定保存model.module.state_dict()还是model.state_dict()要有一个统一策略并且在加载函数里做好前缀兼容。我见过不止一个项目因为保存和加载的代码由两个人写一个存了带module.前缀的权重另一个按原始键名加载结果来回折腾了半天。说实话device相关的问题大多是“知其然不知其所以然”造成的。只要理解了张量和参数各有设备属性、跨设备不能直接运算、DataParallel只是把模型复制到多卡然后拼接输出大部分bug都能靠逻辑推理解决而不是靠试错。这套东西并不难但值得花点时间彻底弄明白因为它是所有Pytorch训练代码的地基。
