资讯详情

资讯详情

PyTorch to(device)全解析:设备迁移原理与高频报错速查

1. 先搞清楚to(device)到底在干什么要论PyTorch里最被低估的操作to(device)绝对排得上号。很多人从入门教程里学到的第一句代码就是model.to(device)但真正问一句“这一步到底做了什么底层发生了什么”能答上来的人其实不多。我见过很多同学在单卡CPU上跑得好好的一上GPU就报Expected all tensors to be on the same device然后一脸懵——明明写了.to(device)为什么还报这个错这篇内容就是想把这个操作彻底掰开揉碎。你不需要多深的理论基础只要写过几行PyTorch代码就行我会从概念讲到底层流程再给出一份可以照着写的封装方案最后把常见的报错整理成速查表。适合刚入门的新手也适合那些用了很久但没时间深究细节的老朋友。1.1 device是什么为什么模型和数据必须待在同一个地方device在PyTorch里的字面意思就是“设备”它决定了一个张量Tensor在硬件层面的存储位置。最常见的两个设备是cpu和cuda其中cuda后面还能跟编号比如cuda:0、cuda:1对应多卡环境里的第几张GPU。从PyTorch 1.12往后mps也被支持了专门用来跑Apple Silicon的GPU用法和cuda几乎一模一样。CPU上的张量存储在系统内存里GPU上的张量存储在显存里两者物理隔离开。PyTorch在底层做的所有运算都要求参与运算的张量在同一个设备上。举个例子import torch x torch.ones(2) # 默认在 cpu y torch.ones(2, devicecuda:0) # 在显卡0上 # 下面这行会直接报错 z x y报错信息就是RuntimeError: Expected all tensors to be on the same device。这有点像把食材放在冰箱里把锅放在炉子上人站在料理台前却没法把两边东西凑到一起——不是菜谱不对而是你做菜时根本没把食材挪到灶台旁边。模型的参数在底层也是Tensor。你用model.to(device)把模型搬到GPU后它内部的权重、偏置、BatchNorm的均值和方差全部同步搬过去。之后跑前向传播输入数据也必须跟着到GPU上否则前向传播时weight input这一下就会因为设备不一致直接炸掉。很多新手困惑的点是为什么报错的不是“模型没搬过去”因为报错信息是运行时才暴露出来的而to(device)本身不会做任何校验。它只是执行了一次设备转换至于你转换之后有没有把返回值接住、有没有把后续的输入也转换PyTorch一概不管。这就引出了下一个问题底层到底是怎么搬的。1.2 to(device)和.cuda()的区别以及真正的底层流程先说结论.cuda()只是.to(device)的一种快捷写法。tensor.cuda()等价于tensor.to(cuda)model.cuda()等价于model.to(cuda)。但.to()提供了更多灵活性比如可以同时指定数据类型、可以动态选择设备所以在工程代码里我更推荐统一用.to(device)。底层流程大致是这样的对于CPU上的Tensorto操作会在目标设备上新建一块内存显存或内存再通过内存拷贝把数据搬过去。如果源设备和目标设备都是显卡涉及的是显存之间的拷贝速度相对较快如果从CPU内存拷贝到显存走的是PCIe总线也就是常说的H2DHost to Device拷贝整个过程是异步的由CUDA流来调度。这里有个容易被忽略的点张量从CPU搬到GPU后返回的是一个新的Tensor对象原来的CPU Tensor不会变。这一点和Module.to()不同——模型的to()是就地修改的参数对象直接换了存储位置而Tensor的to()永远返回新对象。后面我会专门展开这个细节。2. 核心用法拆解一份能直接抄的to(device)手册很多教程把to(device)一笔带过导致很多同学只会一种写法model.to(device)。但其实这个API的完整形态很丰富掌握了才能真正做到“哪里都能跑”。2.1 Tensor的to()设备转换与类型转换一起搞定对一个Tensor调用to()最常见的形式有这么几种import torch x torch.randn(3, 4) # 1. 只指定设备 y1 x.to(cuda) y2 x.to(devicecuda:0) # 2. 只指定数据类型 y3 x.to(torch.float16) # 3. 设备和数据类型同时指定 y4 x.to(devicecuda, dtypetorch.float16) # 4. 可以指定内存格式 y5 x.to(devicecuda, memory_formattorch.channels_last)第1种最常用但有个细节要注意字符串cuda在单卡环境等价于cuda:0但在多卡环境里它会选择哪个编号答案是受CUDA当前上下文影响通常是设备0。如果你的代码里设置了其他设备比如torch.cuda.set_device(2)那么cuda就会默认指向编号为2的卡。为了减少不确定性我在多卡项目里都会显式写cuda:0或cuda:1而不是笼统地写cuda。第2种类型转换其实很实用。比如你从DataLoader里拿出来的图片是float32但推理时想用float16加速一条to(dtype...)就搞定了不需要先.half()再.to()两次调用。有一个我踩过的坑当Tensor同时做设备迁移和类型转换时to的源码会先做类型转换再做设备迁移这在绝大多数场景下没问题。但如果你从cuda上的半精度张量转回CPU的float32记得顺序不影响结果真正影响结果的是你有没有写dtype。我自己曾经只写了.to(cpu)结果从模型输出的float16张量变成了CPU上的半精度Tensor后续数据预处理全跑偏了。2.2 Module的to()模型参数和buffer是重点nn.Module的to()是另一个完全不同体量的存在。它在PyTorch源码里做了递归操作遍历模块的所有子模块、所有参数Parameter、所有buffer然后把它们统一搬过去。这里需要解释一下parameter和buffer的区别。Parameter是用来做梯度更新的权重比如卷积层的weight、bias。Buffer则是模型里需要持久化但不需要梯度的变量最典型的是BatchNorm层里的running_mean和running_var——它们会随着前向传播不断更新但不参与梯度计算。自定义模型时如果你在代码里用self.register_buffer()注册了变量那么model.to(device)也会自动把它搬过去。但如果你只是给模型加了一个普通的Tensor属性比如self.my_tensor torch.zeros(3)那么model.to(device)是不会迁移这个属性的。很多奇奇怪怪的“为什么我的模型在GPU上某个变量还在CPU上”问题根源就在这。最好的验证方式是这样import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 16, 3) self.register_buffer(my_buffer, torch.zeros(3)) self.my_tensor torch.ones(3) # 这行不会被迁移 m MyModel() m.to(cuda:0) print(m.conv.weight.device) # cuda:0 print(m.my_buffer.device) # cuda:0 print(m.my_tensor.device) # cpu注意这里另外模型在GPU上的加载方式也有讲究。我见过有人这么写model MyModel().to(cuda:0) state_dict torch.load(weight.pth) model.load_state_dict(state_dict)如果weight.pth里存的权重原来是CPU上保存的大概率是那么这种写法在PyTorch里是允许的load_state_dict在加载时会自动把state_dict拷贝到与模型参数相同的设备上。但如果反过来模型在CPU上state_dict是从GPU上保存的加载时就会暴露类型或设备不一致的报错。最稳的API是torch.load(path, map_locationcpu)先把权重无条件拉回CPU再加载到模型所在设备这套组合拳我用了很久很少出问题。2.3 非阻塞模式与pin_memory别小看这个参数to()方法还有一个容易忽略的参数non_blocking。x x.to(cuda:0, non_blockingTrue)这个参数的含义是当数据从CPU搬到GPU时让拷贝操作异步执行不用等拷贝完成就让代码继续往下走。听起来很美好能提升性能但它有个前提CPU源Tensor本身必须是固定的内存也就是处于pinned memory。如果你用的是普通内存non_blockingTrue并不会带来性能提升甚至可能因为不安全的异步操作导致结果错误。在PyTorch里pinned memory通常通过DataLoader的pin_memoryTrue参数来生成。当DataLoader启用了pin_memory它产出的Tensor会被放进页锁定内存此时配合.to(device, non_blockingTrue)数据拷贝的H2D部分就能和当前计算流重叠起来相当于一边在算上一批数据一边在搬运下一批数据训练速度会有可见提升。我实测在数据预处理比较重的场景里这个组合能让整体吞吐提升5%-10%左右。不过如果你用的是Windows系统或模型参数本身就很小这个收益会被削弱。我的习惯是pin_memoryTrue无脑开non_blockingTrue在训练循环里写上也没副作用推理时图省事就不加区别不大。3. 实战中的常见翻车现场到了实战阶段问题变得五花八门。这里我把近几年见过最多的几个翻车现场整理出来每一个都是真实发生过、真金白银踩过坑的。3.1 最容易犯的错忘了给to()重新赋值这是Tensor的to()和Module的to()行为不一致导致的最经典事故。# 错误示范 x torch.randn(4, 3) x.to(cuda:0) # 返回值是新的Tensor但没接住 print(x.device) # cpu原Tensor完全没变 # 正确示范 x x.to(cuda:0) # 重新赋值因为PyTorch的张量操作总体遵循“不修改原对象”的原则to()返回新张量原地版本并不存在。很多新手把Tensor当成list去用觉得调一下方法就变了结果后续一运算就报设备错误。Model的to()恰好相反它是就地操作的不需要赋值。但即便如此我也见过有人在自定义的训练框架里写出model model.to(device)——多了个赋值其实也没问题因为它返回的就是同一个对象只是很多余。更隐蔽的变种是在Dataset里对数据进行预处理后忘了把结果返回出去或者在损失函数里手动构造了一个Tensor时忘了指定device。这种问题排查起来特别耗时间因为报错往往不是出现在你漏写的那一行而是出现在几百行之后真正的计算点。我的建议是在训练脚本最开始就用一个统一的device变量所有Tensor创建的地方都写torch.zeros(..., devicedevice)而不是创建之后再调.to()。第二种方式虽然也能用但容易出现漏网之鱼。3.2 多GPU/分布式场景下的特殊用法多卡环境下device的选择需要更讲究一些。最常见的是DataParallel它的写法是model nn.DataParallel(model) model model.to(cuda:0)DataParallel内部会把输入张量在batch维度上前向切分给多张卡那么输入数据到底要放在哪张卡上答案通常是主卡也就是默认的cuda:0。所以书写方式依然是先把输入搬到cuda:0剩下的交给DataParallel内部处理。另一种更推荐的方式是DistributedDataParallel。它要求每个进程绑定一张卡通常用local_rank来指定import os local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) device fcuda:{local_rank} model model.to(device)这里的device是每个进程各自不同的值数据也搬到自己的卡上。很多新手在分布式脚本里直接用devicecuda:0硬写结果每个进程都把数据往同一张卡上搬轻则性能极差重则OOM。正确做法永远是用本地rank动态构造device。还有CUDA_VISIBLE_DEVICES这个环境变量。它会影响程序能看到的显卡列表设置之后原来的物理卡号会被重新映射。比如物理上有4张卡你设置CUDA_VISIBLE_DEVICES2,3那么程序里cuda:0对应的就是物理上的第3张卡。这种映射关系对to(device)是透明的但如果你在代码里通过torch.cuda.get_device_name(0)去看显卡名字就会发现跟你预期的不一样。3.3 混合精度训练中的device坑混合精度训练用上AMPAutomatic Mixed Precision之后to(device)的场景会多一层维度不仅要管设备还要管数据类型。AMP的标准用法是把模型和数据都放在GPU上但用torch.cuda.amp.autocast()包裹前向传播让PyTorch自动选择合适的精度。这时候有个容易犯的错在autocast区域外手动把激活值切成float16然后在区域内计算。这种写法本身不报错但会让某些OPS跑在低精度上造成精度不稳定。还有一种常见操作是把模型直接变成半精度模型通过model.half()实现。这里要注意顺序如果你先执行model.half()再执行model.to(device)那模型参数就会保留为float16并搬到GPU如果先.to(device)再.half()也一样参数从float32变成float16后还在原来的设备。两者结果一致但有个中间态容易踩坑# 有问题的写法先half()此时参数还是float16但没搬到GPU model.half() loss compute_loss(model, data) # data在cpu模型也在cpu # 后续再model.to(device)也能用但缓存了cpu上的半精度权重最规范的是只在最后推理阶段使用half()配合to(device)训练阶段交给AMP自动管理。毕竟手动全半精度训练损失函数和梯度更新那边很容易出现精度溢出我吃过几次亏之后就再也不手动half了。4. 直接可用的封装方案与调试技巧写到这里你应该对to(device)有了完整的理解。接下来分享一些我在实际工程里沉淀下来的封装和调试手段这些属于常规文档里不会写的内容。4.1 封装一个统一的device工具函数项目刚开始就先把设备管理工具写好能省后续一多半的烦恼。我的做法是单独建一个device_manager.py全局只用一个函数import torch def get_default_device(): if torch.cuda.is_available(): return torch.device(cuda) elif hasattr(torch.backends, mps) and torch.backends.mps.is_available(): return torch.device(mps) return torch.device(cpu) device get_default_device()然后在所有需要构建Tensor的地方都直接传入这个全局deviceinput_tensor torch.randn(4, 3, devicedevice)有人可能会问那mps会不会在写cuda的地方报错会的。所以统一封装后你的代码里就不应该再出现裸写的cuda字符串除非你明确只支持N卡。另一个常用封装是safe迁移函数def safe_to(x, device): if x.device device: return x return x.to(device)这个函数能避免重复的拷贝操作。尤其在一个值可能已经被搬到目标设备的情况下直接用.to()也没错但会多一次判断逻辑之外的内存分配。在训练循环里频繁调用时多一次不必要的copy可能有肉眼可见的性能损耗。4.2 在代码里验证到底跑在哪个设备上调试设备问题最直接的手段就是在可疑位置打印Tensor.device、Parameter.device。print(finput device: {input_tensor.device}) print(fmodel weight device: {next(model.parameters()).device})更工程化的做法是加断言让错误在最早的时间点暴露assert input_tensor.device device, f输入张量不在目标设备上期望{device}实际{input_tensor.device}这种断言放在训练循环开头即可成本极低。我经历过很多次“训练了10个epoch后才在评估阶段发现某些统计量一直在CPU上”的事故自从加了断言后这类问题基本当天就暴露。另外一个有用的trick在DataLoader的collate_fn里完成设备迁移让数据一出来就在GPU上这样后续代码里就不用到处写.to(device)了。代码长这样def collate_fn(batch, device): images torch.stack([item[image] for item in batch]).to(device) labels torch.tensor([item[label] for item in batch]).to(device) return items4.3 关于显存占用和内存拷贝的个人经验最后说几个我不太容易在官方文档里找到、但实际工程中反复验证过的经验。第一显存碎片的累积问题。如果你的训练循环里频繁出现大张量.to(device)和.cpu()切换显存会出现碎片化。表现是nvidia-smi显示显存还有空间但一分配大Tensor就OOM。解决方案是尽量避免在循环内来回搬运数据尽量让数据在一个设备上待久一点需要聚合统计的时候再搬回CPU。第二反向传播之前确认 loss 也在对的设备上。loss.to(device)这种写法我见过很多次但它毫无意义——loss就是模型输出直接算出来的它在哪个设备上取决于模型和在哪个设备上。你真正要关注的是label有没有搬到GPU而不是loss。第三梯度清零不会影响设备。optimizer.zero_grad()只是把梯度Tensor的数值置零不会把Tensor搬到CPU。这句话听起来是废话但我真遇到过有同学以为.zero_grad()会把梯度数据移走于是在每个step结束都加一句model.to(device)。第四如果你要保存模型最好在CPU上做个统一入口。我的习惯是保存前先model.cpu()然后torch.save(model.state_dict(), path)。这样保存出来的权重文件里全是的CPU Tensor换任何机器加载都方便再用map_location加载到目标设备一劳永逸。5. 常见报错与排查速查表把这些年遇到最多的问题整理成一张表方便你以后碰到报错直接对照。报错或现象根本原因解决办法Expected all tensors to be on the same device参与运算的张量不在同一设备用.to(device)统一输入输出检查漏网Tensor打印device定位Torch not compiled with CUDA enabled安装的PyTorch是CPU版卸载后安装CUDA版或检查虚拟环境是否被全局覆盖模型前向传播正常但反向传播报错部分自定义Tensor没有被Module.to()迁移使用register_buffer注册或手动迁移该属性保存权重后再加载报错size mismatch模型结构和权重状态不一致不一定是device问题先检查state_dict的键名和shape设置CUDA_VISIBLE_DEVICES后程序仍然OOM设备编号被系统重新映射在脚本开头打印可见设备列表用torch.cuda.device_count()确认使用了pin_memoryTrue但non_blockingTrue无效果CPU源数据不是固定内存确认DataLoaderpin_memoryTrue如果是自定义数据集手动分配固定内存GPU上程序跑通但速度比CPU还慢单次数据量太小H2D拷贝开销太大批量处理或把数据生成部分尽量放到GPU上DataLoader迭代时数据一直在CPU上collate_fn里没做迁移在collate_fn或训练循环开头统一迁移多进程训练时每张卡都报OOM所有进程都往同一张卡上搬数据用local_rank动态构造cuda:{local_rank}不要硬编码卡号模型参数和输入都在GPU但loss突然变成NaN半精度溢出常见于fp16检查是否手动.half()了不该半精度的部分改用AMP的autocast这张表并不能覆盖所有场景但覆盖了我在社区答疑里见到的高频问题。碰到没列出来的先打印所有相关Tensor的.device和.dtype大部分问题一眼就能定位。最后再分享一个我自己的习惯在写模型框架的时候把device作为超参数传入而不是在模型内部硬编码cuda。这样同样的代码能跑在CPU、单卡、多卡、甚至是MPS设备上切换设备只需要改一个配置项。这个习惯帮我省了太多来回改代码的时间也让迁移到新环境时最大的障碍从改代码变成了装驱动。希望这篇内容能帮你把to(device)这块硬骨头啃下来。下次再遇到设备不一致的报错先别急着改代码想想那个Tensor它到底被安排到哪个设备上了——大多数时候问题出在“你以为是A设备其实是B设备”这一念之间。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →