资讯详情

资讯详情

PyTorch Tensor设备管理与GPU加速:搬运、显存与OOM排查

拿 Tensor 往 GPU 上搬这件事看着就是一行.to(cuda)但我见过太多人卡在这行之后的一堆问题上报错Expected all tensors to be on the same device、显存莫名其妙爆掉、明明用了 GPU 却比 CPU 还慢、多卡环境下程序抢错卡。这篇就聊 PyTorch 里最基础也最容易翻车的一环——Tensor 的设备管理与 GPU 加速。内容从张量在内存里的真实形态讲起到 CPU/GPU 之间的数据搬运、显存占用估算、OOM 排查路径再到我自己实际踩过的坑全部配可运行的代码和实测思路。适合刚上手 PyTorch、正准备把训练脚本从 CPU 迁到 GPU 的同学也适合已经能跑通但总感觉“没跑满”的读者——很多性能问题的根因其实就在设备管理这几行代码里。1. 张量搬上 GPU 之前先搞清楚它在搬什么1.1 Tensor 的本质是一块带类型注解的连续内存很多人把 Tensor 当成“多维数组”就完事了这个理解在 CPU 上够用一上 GPU 就不够。准确说一个 Tensor 由三部分组成一块存储storage就是一段连续的、按字节计的内存一组元数据metadata包括形状、步长、偏移量以及两个属性——**dtype数据类型**和device所在设备。形状和步长只是“怎么解读这块内存”的说明书真正占地方的是那块存储。这个区分非常重要因为它解释了一堆看起来反直觉的现象。比如你用transpose或者permute转置一个张量PyTorch 几乎不花时间因为它只改了步长元数据底层存储一个字节都没动。再比如切片x[0:2]得到的新张量跟原张量共享同一块存储你改其中一个另一个也会变。同理view要求内存连续才能用reshape在不能view的时候会悄悄复制一份——这一复制显存就多了。device这个属性决定了这块存储到底落在哪。CPU 张量的存储是主机内存host memoryGPU 张量的存储是显存device memory。两者之间没有共享地址空间所以所谓“搬”本质是一次跨总线的内存拷贝。理解了这一点后面所有的性能问题都有了统一的解释框架跨设备拷贝永远比设备内的计算贵得多能少搬就少搬。1.2 CPU 和 GPU 的差距不在主频在算术密度经常有人问为什么同一段矩阵乘法GPU 能快几十倍甚至上百倍答案不是“GPU 核心多”这么简单而是算术强度这个指标。算术强度 浮点运算次数 / 访存字节数单位是 FLOP/Byte。拿方阵乘法C A B矩阵规模 N×Nfloat32 存储。计算量是 2N³ 次浮点运算访存量大约是 3N² 个元素乘以 4 字节。两者一除算术强度约为 N/6 FLOP/Byte。代入数字感受一下。N64 时算术强度只有 10 左右N1024 时大概 170N4096 时接近 683。而 CPU 平台的内存带宽通常在几十 GB/sGPU 的显存带宽则是数百 GB/s 到 TB/s 量级。乘以算术强度就能看出在 N 足够大的时候GPU 的算力上限被拉高了整整一个量级。反过来推这也直接给出了不该上 GPU 的场景N 很小的时候比如 32、64 这种算术强度低你会彻底卡在访存上GPU 的算力根本喂不饱加上核函数启动的固定开销几微秒到几十微秒实际耗时可能还不如 CPU。我在做序列标注那会儿有个逐样本的小矩阵运算循环里每次都调一次 GPU结果比纯 CPU 版本慢了将近 3 倍改成先把一批样本拼成一个大张量再一次性送进去才把速度拿回来。1.3 判断标准这三类操作值得搬到 GPU基于上面的逻辑我一般用三条经验判断第一计算密度高的批量操作。矩阵乘、卷积、大批量的逐元素运算这些是 GPU 的主场规模越大越明显。第二同一份数据要被反复使用。典型的就是模型参数——它在整个训练过程中一直待在显存里每次前向反向都复用搬运成本被摊薄到可以忽略。第三批大小能撑起来。这点常被忽视。batch size8 的小卷积GPU 利用率可能只有百分之十几同样的模型换成 batch size256吞吐能翻好几倍。如果确实受限于显存或者其他原因只能用小 batch那就要靠累积梯度或者提高数据加载效率来补而不是硬指望 GPU 自动变快。至于逐元素的细粒度小张量、需要在 Python 层频繁做条件判断的逻辑、以及循环内的小规模运算老老实实放 CPU或者先做批量合并再送 GPU。这条边界认清楚了后面的所有优化才有意义。2. 设备对象与环境检查把 device 当一等公民2.1 起手第一件事把设备信息打印清楚我见过不少人拿到别人代码直接跑报一堆设备相关的错其实就是环境没对齐。花两分钟做一次环境体检能省后面两小时。下面这段我几乎每个新环境都会跑一遍。import torch print(torch version:, torch.__version__) print(cuda available:, torch.cuda.is_available()) print(cuda version:, torch.version.cuda) print(device count:, torch.cuda.device_count()) if torch.cuda.is_available(): idx torch.cuda.current_device() prop torch.cuda.get_device_properties(idx) print(device name:, prop.name) print(total memory: %.2f GB % (prop.total_memory / 1024**3)) print(multi processor count:, prop.multi_processor_count) print(capability:, torch.cuda.get_device_capability(idx))这里几个字段各有用途。is_available()为 False通常是驱动或者 CUDA 运行时没配好也可能是装成了 CPU-only 版本的包——这种情况特别隐蔽因为import torch一点问题没有只是永远用不上 GPU。torch.version.cuda是编译时锚定的 CUDA 版本它和系统里驱动支持的版本之间有兼容关系版本对不上会直接报错。total_memory决定你能开多大 batchmulti_processor_count这些是粗看算力档位capability计算能力形如 8.6、8.9 这种则决定了能不能用上某些特定指令集——比如半精度相关的加速单元在老卡上是不支持的代码里开了 autocast 反而可能更慢。提示torch.cuda.get_device_properties返回的显存大小是物理显存实际可用会略少一些因为运行时本身要占用一部分上下文开销。2.2 统一用 device 变量不要在代码里写死 cuda这是我认为最值得养成的一个习惯在脚本入口处确定一次设备之后全程传递这个 device 对象。device torch.device(cuda if torch.cuda.is_available() else cpu) # 后续所有张量创建、模型迁移统一走这个变量 x torch.randn(4, 8, devicedevice) model MyModel().to(device)对比一下常见的两种写法。一种是到处写x.cuda()这种做法的问题是没有 GPU 的机器上直接崩别人拿到你的代码第一件事就是全局搜索替换。另一种是写死torch.device(cuda:0)单卡机器没问题一到多卡环境就容易出岔子。用统一的 device 变量还有个隐性好处代码路径唯一。CPU 和 GPU 走同一套逻辑只是 device 不同调试的时候只要把 device 切成 CPU 就能排除掉一大类设备相关的干扰因素。我在排查数值不一致问题的时候这个技巧用过很多次——先切 CPU 跑一遍结果如果正常那问题基本锁定在 GPU 特有的路径上比如半精度、某些核函数的实现差异、非确定性算法。2.3 多卡环境下的设备可见性控制多卡机器上有个经典问题所有人都默认抢 0 号卡其他卡闲着。控制可见设备有两个层面。第一个是环境变量层面在进程启动前设定CUDA_VISIBLE_DEVICES2,3 python train.py设完之后进程内看到的逻辑编号会重新映射——物理上的 2 号卡变成cuda:03 号变成cuda:1。这个映射关系一定要记住否则日志里打印的cuda:0和你以为的物理卡号会对不上排查显存问题时非常误导人。第二个是代码层面用torch.cuda.set_device(idx)指定当前线程的默认设备。这个设置是线程局部的多线程或者多进程的场景下每个进程要各自设一次。import os import torch local_rank int(os.environ.get(LOCAL_RANK, 0)) torch.cuda.set_device(local_rank) device torch.device(cuda, local_rank)注意set_device只影响“不指定索引时默认用哪张卡”它不会阻止你显式地往别的卡上放张量。真正做隔离还得靠可见设备列表。我在一个共享集群上就因为这点吃过亏——以为设了默认设备就安全了结果代码里某处遗留的.cuda()把张量扔到了别的卡上训练跑了半小时才发现显存统计完全对不上。3. 跨设备搬运to、cuda、cpu 到底该用哪个3.1 三种写法的实际差别把张量在设备间搬运PyTorch 提供了好几个入口.to(device)、.cuda()、.cpu()还有torch.Tensor.to的 dtype 重载。功能上有重叠但用法上有讲究。.to(device)是通用写法能同时处理设备和类型也能接受torch.device对象是唯一推荐在正式代码里用的。.cuda()是它的历史遗留快捷方式等价于.to(cuda)好处是短坏处是硬编码了后端。.cpu()反向操作功能没问题但如果代码要同时兼容其他加速后端就得改成.to(cpu)保持风格统一。有一个细节值得说清楚这些方法返回的都是新张量原张量不变。因为 Tensor 的 device 属性不可原地修改——你不能对一个已经在 CPU 上的张量执行x.device cuda会直接抛AttributeError。所以正确的写法永远是x x.to(device)或者干脆在创建时就指定device。后者更好因为省掉了一次中间张量的创建也就省了一份临时内存。# 不推荐先建再搬多一份临时内存 x torch.randn(1000, 1000) x x.to(device) # 推荐创建时就落在目标设备上 x torch.randn(1000, 1000, devicedevice)别小看这个差别。1000×1000的 float32 张量是 4MB看着不多但如果在数据处理循环里这么写几千次下来分配器会被临时张量搅得很碎显存碎片的积累会比你想象的快。3.2 固定内存与非阻塞拷贝的配合这是本节最有价值的部分。默认情况下从 CPU 页内存往 GPU 拷贝是同步的CPU 发出拷贝指令后必须等拷贝完成才能继续。而如果源内存是页锁定内存pinned memory拷贝就可以走 DMA 直接传输并且支持异步。# 普通拷贝同步CPU 会阻塞等待 x torch.randn(1_000_000) x_gpu x.to(device) # 固定内存 非阻塞异步可与其他计算重叠 x torch.randn(1_000_000).pin_memory() x_gpu x.to(device, non_blockingTrue)pin_memory()的作用是把这块内存锁在物理页里不被操作系统换出。好处是 GPU 的 DMA 引擎可以直接读不用先做一次中间拷贝代价是这块内存不能再被换页分配和释放都比普通内存慢占用也会让系统的可用物理内存变少。那收益到底有多大做个粗略测算。假设你在训练循环里每个 step 都要传 100MB 数据PCIe 4.0 x16 的有效带宽实测大概在 20~25 GB/s理论传输时间约 4~5ms。同步拷贝的话这 5ms CPU 完全干等着非阻塞加双缓冲的话这 5ms 可以和上一批的 GPU 计算重叠掉。如果单个 step 的 GPU 计算时间是 30ms那省下的 5ms 就是大约 14% 的端到端加速。数据量越大、计算越重这个收益越明显。组合方式CPU 是否阻塞能否与计算重叠适用场景普通内存 同步拷贝是否一次性小数据传输普通内存 non_blocking名义异步实际仍同步否无实际意义不推荐固定内存 同步拷贝是否传输本身更快但仍阻塞固定内存 non_blocking否是数据加载流水线、大批量传输注意表格第二行non_blockingTrue只有在源内存是 pinned 的前提下才真正生效普通页内存传参会退回同步行为。这个坑很隐蔽——你写了non_blockingTrue代码不报错但一点加速都没有。3.3 用 DataLoader 把搬运藏进流水线真实训练里数据加载的搬运不该由你手动管交给DataLoader更省事from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue, prefetch_factor2, persistent_workersTrue, ) for batch in loader: inputs batch[input].to(device, non_blockingTrue) labels batch[label].to(device, non_blockingTrue) # 前向、反向这里几个参数的作用值得说透。pin_memoryTrue让 worker 进程把产出的 batch 放进固定内存池这是后面非阻塞传输能生效的前提。num_workers是子进程数量取值经验是 CPU 核数的 1/4 到 1/2 之间太少喂不饱 GPU太多会争抢 CPU 和内存带宽反而拖慢。prefetch_factor控制每个 worker 预取多少个 batch默认是 2显存宽裕、计算重的时候可以适当调高。persistent_workersTrue是我强推的一个参数。它让 worker 进程在 epoch 之间保持存活避免每个 epoch 结束都重新 fork 一次。数据量小、epoch 多的时候这个开销能占到总时间的相当比例——我实测过一个数据集平均每个 epoch 训练 8 秒但 worker 重建要花 1.5 秒左右加了这个参数之后直接降到 8.2 秒以内。提示设了pin_memoryTrue之后可以用torch.cuda.memory_stats()观察pinned_bytes相关指标但更直接的方法是看 GPU 利用率曲线。如果利用率呈锯齿状周期性掉到 0说明数据供给断流了优先调num_workers和prefetch_factor。4. GPU 上的核心操作与自动求导联动4.1 创建、变形、索引在 GPU 上的行为差异张量到了 GPU 上绝大多数操作的行为和 CPU 完全一致但有几个地方要注意。创建类操作要显式指定设备否则默认落在 CPU 上。torch.zeros、torch.ones、torch.arange都支持device参数一定要带上。像torch.tensor([1, 2, 3])这种从 Python 列表构造的哪怕你后面.to(device)中间也已经在 CPU 上分配过一次了。变形类操作里view要求内存连续如果张量是transpose之后的结果view会报错。这时候要么用contiguous()先复制一份多占一份显存要么用reshape内部自动判断必要时复制。在显存紧张的场景下这个隐性复制会咬人——我见过有人在一层里连续做了四次reshape且都不是 view凭空多出好几百 MB 占用。索引类操作有个硬性要求用作索引的张量必须是longint64类型而且最好和待索引张量在同一设备上。跨设备的索引在某些版本里能跑但会触发隐式同步拷贝性能很差。x torch.randn(1000, 128, devicedevice) idx torch.tensor([3, 7, 11], dtypetorch.long, devicedevice) # 正确dtype 和 device 都对 selected x[idx] # 错误示范int32 索引直接报错 bad_idx torch.tensor([3, 7, 11], dtypetorch.int32, devicedevice) # selected x[bad_idx] # 抛 IndexError4.2 dtype 与 device 的联动陷阱半精度float16 / bfloat16是 GPU 加速里绕不开的话题但它和 device 的组合有不少坑。第一不是所有 GPU 都支持高效的半精度运算。前面提到的计算能力字段在这里起作用较老的架构上 float16 运算会被拆成多个单精度指令模拟速度反而更慢。所以上精度优化之前先确认硬件支持。第二float16 的数值范围窄。它的最大可表示值大约 65504最小正规数约 6e-5训练中梯度很容易下溢成 0 或者上溢成 inf。这也是为什么通常要配GradScaler做损失缩放——把 loss 人为放大一个系数让梯度落在可表示范围内反向传播完再缩回来。from torch.amp import autocast, GradScaler scaler GradScaler() for inputs, labels in loader: inputs inputs.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) optimizer.zero_grad(set_to_noneTrue) with autocast(device_typecuda, dtypetorch.float16): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()第三混合精度下的设备一致性。autocast 会自动把部分算子降到半精度但模型参数本身可能仍是 float32。这种混合状态是设计使然不用手动干预但如果手写自定义算子就要自己处理类型转换否则会报expected scalar type相关的错误。还有一个容易被忽略的点float16 的张量做归约sum、mean时中间累加最好用 float32。对一个长度几万的向量求和在 float16 下做累积误差会很明显甚至可能因为溢出得到 inf。PyTorch 的部分归约算子内部已经做了类型提升但不是所有算子都这么贴心涉及数值敏感的计算时手动.float()一下更稳妥。4.3 反向传播中的设备一致性自动求导对设备的要求很严格计算图上的所有张量必须在同一设备上。前向时漏了一个张量没搬反向传播时会在那个节点炸掉报错信息通常指向Expected all tensors to be on the same device, but found at least two devices。叶子和非叶子张量的梯度存放位置有规律loss.backward()之后非叶子张量的.grad默认是 None不保留叶子张量的.grad会落在和该叶子张量相同的设备上。所以如果模型已经在 GPU 上参数的梯度自然也在 GPU 上这点不用担心。但如果你手动创建了一个需要求导的常量张量在 CPU 上再拿它去和 GPU 张量运算就会触发前面说的报错。w torch.randn(8, 8, devicedevice, requires_gradTrue) b torch.randn(8, 8) # 忘了指定 device落在 CPU x torch.randn(4, 8, devicedevice) try: y (x w b).sum() y.backward() except RuntimeError as e: print(报错, e)这类问题在重构代码时特别容易引入——原来全在一个文件里设备清清楚楚拆成多个模块之后某个模块内部悄悄建了个 CPU 张量就要花时间找了。我的做法是在每个模块的入口加一个断言def forward(self, x): assert x.device self.weight.device, f设备不一致: {x.device} vs {self.weight.device} ...多写一行断言少查两小时 bug这笔账很划算。注意requires_gradTrue的张量只能在浮点类型上创建。整型张量设了会直接报错only Tensors of floating point and complex dtype can require gradients。索引、掩码这类张量天然不需要梯度保持整型或布尔型就对了还能省显存。5. 显存管理比算力更容易先撞上的墙5.1 先把显存占用算清楚OOM 是最常见的 GPU 问题但大部分人只会盲目调小 batch size。真正有效的做法是先把占用拆开看看每一块占多少。训练时的显存占用大致分四块占用项估算方式float32说明模型参数参数量 × 4 字节常驻训练和推理都有梯度参数量 × 4 字节仅训练时存在优化器状态参数量 × 4 或 × 8 字节Adam 是两倍SGD 无激活值与 batch、序列长度、层数正相关最大头也最不可控临时缓冲视算子而定cuBLAS 工作区、归约中间量拿一个 1.1 亿参数的模型举例。参数约 440MB梯度 440MBAdam 的两个状态各 440MB光这三项就是 1.76GB。听着不多但激活值往往才是大头——同样是这个模型batch size 32、序列长度 512 的情况下激活值可能占到好几 GB具体取决于层结构和是否用了梯度检查点。算这笔账的意义在于知道哪块是固定成本哪块能压缩。参数、梯度、优化器状态基本是常数除非换优化器或者做分布式切分激活值是唯一能通过调整 batch size 直接控制的部分。所以看到 OOM 先调 batch逻辑是对的但如果发现 batch 降到 1 还在 OOM那就得看参数加载和临时缓冲了。5.2 用内存统计定位问题PyTorch 提供了一组内存查询接口很多人不知道它们的存在其实非常好用。def report_memory(tag): allocated torch.cuda.memory_allocated() / 1024**2 reserved torch.cuda.memory_reserved() / 1024**2 max_alloc torch.cuda.max_memory_allocated() / 1024**2 print(f[{tag}] allocated{allocated:.1f}MB reserved{reserved:.1f}MB peak{max_alloc:.1f}MB) report_memory(开始) # ... 某些操作 ... report_memory(加载模型后)这里allocated和reserved的区别是关键知识点。allocated是实际被张量占用的显存reserved是缓存分配器向驱动申请下来的总额度。分配器从驱动一次性申请一大块之后在内部切分复用所以reserved通常大于等于allocated差值就是空闲的缓存块。这解释了两件事。其一empty_cache()释放的是缓存分配器里那些没人用的块如果所有块都还被张量引用着调用它一点用没有——所以它不是解决 OOM 的万能钥匙。其二模型训练时显存曲线通常是阶梯式上升的因为每次需要更多空间时分配器会整块申请而不是按需逐个申请看起来像突然跳了一大截别慌那是缓存。torch.cuda.memory_summary()能打出更详细的快照包括不同尺寸段的块分布。判断是否存在碎片问题就看有没有大量小块空闲但拼不出大块连续空间——这种情况下reserved很接近物理上限而allocated却不高。5.3 显存优化的四个层次按见效快慢和侵入性排序我一般这么处理第一层改 batch size 和累积梯度。最直接。显存不够就减 batch减到不影响收敛的下限后用梯度累积凑回等效批量。要点是 loss 要除以累积步数否则等效学习率会翻倍。accum_steps 4 for i, (inputs, labels) in enumerate(loader): inputs inputs.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) with autocast(device_typecuda, dtypetorch.float16): loss criterion(model(inputs), labels) / accum_steps scaler.scale(loss).backward() if (i 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_noneTrue)注意zero_grad(set_to_noneTrue)这个细节。默认行为是把.grad置零这需要遍历所有参数并写一遍显存设成True则直接把.grad设为 None反向传播时重新分配省掉了一次写内存的操作实测能有几个百分点的加速。第二层推理不开计算图。只做推理的时候一定要包torch.no_grad()或者torch.inference_mode()。前者会让计算图上的节点不再记录后者更彻底连版本计数器都不维护。同一个模型推理时忘了加这个显存占用可能翻好几倍——因为中间激活全被保留下来等着算梯度而实际上根本不会算。model.eval() with torch.inference_mode(): for inputs in test_loader: inputs inputs.to(device, non_blockingTrue) logits model(inputs)第三层及时解引用。在大循环里如果某个中间张量用完了就del掉再配合偶尔调一次empty_cache()能缓解碎片。但注意del只减少引用计数如果还有别的变量指着同一块存储内存不会释放。del intermediate torch.cuda.empty_cache()empty_cache()别放在循环里频繁调。它本身有同步开销每次都会让 CUDA 上下文做一次整理密集训练时可能拖慢整体速度。我的习惯是在每个 epoch 结束时或者加载新模型前调一次。第四层梯度检查点checkpointing。用计算换显存思路是不保存中间激活反向传播时重新算一遍前向。代价是训练变慢大概 20%~30%收益是激活值占用大幅下降能让 batch size 放大好几倍往往整体吞吐还是提升的。from torch.utils.checkpoint import checkpoint def segment_forward(x): return checkpoint(self.block, x, use_reentrantFalse) h segment_forward(x)use_reentrantFalse这个参数在新版本里建议显式指定它对应的是非重入的实现兼容性更好也支持一些重入版本不支持的场景比如输入里有需要求导的张量、或者和 RNG 状态相关。6. 常见问题速查与踩坑实录6.1 报错信息对照表报错关键字根因处理方式Expected all tensors to be on the same device运算双方的 device 不同检查各输入和参数的 device逐层加断言定位CUDA out of memory显存不足先降 batch再看是否漏了 no_grad 或用memory_summary查碎片only Tensors of floating point ... can require gradients对整型张量设了 requires_grad改成浮点类型或去掉该设置Expected object of scalar type Long but got ...索引张量 dtype 不对转成torch.longcuda runtime error: device-side assert triggered越界索引、非法值常出现在CrossEntropy的目标越界先切 CPU 复现定位CUDA error: an illegal memory access核函数访问越界或异步错误先同步一次定位到具体算子再查张量形状Cannot re-initialize CUDA in forked subprocess多进程里在 fork 之前碰了 CUDA把 CUDA 初始化挪到子进程内部关于最后一条补充一句用DataLoader的 worker 时如果 dataset 的__init__里就创建了 CUDA 张量多进程启动会触发这个错误。正确做法是 worker 里只用 CPU产出 batch 之后再在训练进程里搬。6.2 性能异常的三步定位法如果 GPU 跑起来但感觉不对我按这个顺序查第一步看 GPU 利用率。用系统监控工具观察如果长时间低于 50%说明 GPU 在等数据问题在数据加载链路。优先加num_workers、开pin_memory、提高prefetch_factor。第二步看显存和计算的时间比。在训练循环里插计时代码分别记录数据搬运、前向、反向、优化器更新的耗时。import time torch.cuda.synchronize() t0 time.perf_counter() # 待测代码段 # ... torch.cuda.synchronize() t1 time.perf_counter() print(f耗时 {(t1 - t0) * 1000:.2f} ms)这里torch.cuda.synchronize()是必须的。GPU 操作默认异步提交到流里就返回不 sync 的话测出来的是“提交指令的时间”不是真实计算时间结果会小得离谱。很多人测出来发现自己代码快得不可思议就是这个原因。第三步看是否踩了同步点。一些看着无害的操作会强制同步把异步流水线打断.item()、.cpu()、print(tensor)、tensor.tolist()、用张量做布尔判断if tensor:。这些在循环里出现性能会很惨。我见过有人在每个 batch 里print(loss.item())打日志本来能跑满的卡利用率掉到 40% 以下改成攒 N 个 step 打一次就好多了。6.3 几个我踩过的坑坑一.to(device)写在错误的位置。有个项目里我把x.to(device)写在了for循环外面但x在循环里被重新赋值成了 CPU 张量结果从第二轮开始全在 CPU 上跑而且因为 PyTorch 在某些场景下会自动做设备隐式转换部分算子支持代码没报错只是慢。查了两天才发现。教训是搬运的代码要贴着它的消费点写别隔太远。坑二模型初始化后忘了搬。model MyModel()之后直接optimizer Adam(model.parameters())然后训练时把输入搬到了 GPU参数还在 CPU。这时候要么报设备不一致要么在某些自动处理的情况下悄悄做单元素级别的隐式转换慢得离谱。正确的顺序是建模型 →.to(device)→ 建优化器。优化器持有的是参数的引用模型搬了参数对象还是同一批所以顺序其实不严格但习惯上按这个来更不容易错。坑三显存没爆但分配失败。有过一次memory_allocated显示只用了 60%但申请一个大张量时失败。原因是碎片——缓存里都是小块拼不出连续的大块。解决办法是把可能同时存在的大张量提前分配好或者调整分配策略。这类问题在长时间运行、频繁分配释放不同尺寸张量的场景里比较常见。坑四多进程下的设备编号。分布式训练时每个进程要通过LOCAL_RANK决定用哪张卡但LOCAL_RANK的语义是节点内的相对序号不是全局序号。跨节点的时候全局序号得靠RANK算。这点搞混了会出现多个进程抢同一张卡表现是显存莫名翻倍或者直接 OOM。坑五把 CUDA 张量塞进numpy()。tensor.numpy()对 GPU 张量会直接报错必须先.cpu()。而且在 GPU 上.cpu()是同步阻塞的如果循环里频繁这么做性能会掉崖式下滑。数据要交给 CPU 端的库处理时尽量批量搬别一条一条来。提示如果你在调试设备相关的问题一个很实用的办法是在关键位置打印torch.cuda.memory_allocated()和张量的device配合torch.cuda.synchronize()把异步的报错变成同步的即时崩溃堆栈信息会准确得多。异步报错最烦人的点就是真正出错的位置和你看到崩溃的位置可能差了好几步。我个人在实际操作中的体会是GPU 加速这件事写对代码只是及格线能不能跑出应有的性能八成取决于你有没有把它当成一条流水线来看——数据在 CPU 侧准备、经总线过桥、在 GPU 上计算、结果再回来每一段都有延迟每一段都可能成为瓶颈。与其盯着那个.to(cuda)反复怀疑不如在循环里插几次synchronize加计时把每一段的耗时摊开看。数据摆在那里瓶颈在哪一目了然比凭感觉调参高效得多。
觉得有用,分享给同行:

为您的企业打造数字门面

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

立即咨询 →