我第一次把训练脚本从 CPU 挪到显卡上跑是很多年前的事了。当时满心以为改一行.cuda()就能看到几十倍的提速结果跑完一个 epoch计时器上的数字比 CPU 版本还难看。后来花了整整一个下午才发现问题根本不在显卡而在我根本没搞清楚 PyTorch 里 Tensor张量和 GPU 加速之间的关系哪些操作真的被搬到显卡上执行了哪些数据还老老实实躺在内存里被反复搬运哪些地方因为一个小小的设备不匹配在偷偷做隐式拷贝。这件事之后我养成了一个习惯——凡是涉及device的代码一律当成一等公民来对待而不是随手加个.to(cuda)了事。这篇内容聊的就是 PyTorch 数据结构里最基础也最容易被轻视的一环Tensor 的 GPU 加速。它不是什么高深话题但真正把它用明白的人并不多。适合刚接触 PyTorch、想把模型跑到显卡上的新手也适合已经能跑通训练、但总是被显存溢出和设备报错折腾的中级使用者。我会从“为什么快”讲到“怎么搬”再讲到“显存这本账怎么算”和“精度怎么配”最后给一段可复现的基准测试写法。全程都是我自己踩过坑之后总结的东西能直接抄去用。1. 先把“为什么快”这笔账算清楚1.1 显卡不是“更快的 CPU”而是“宽得多的 CPU”很多人对 GPU 加速的第一印象是“显卡主频高”。恰恰相反一张中端显卡的单核频率往往还不如同时代的桌面 CPU。它快的地方在别处CPU 可能只有十几个核心而显卡上有几千个计算单元。你可以把 CPU 想象成一个十几个人的精英小队每个人都能独当一面、能处理复杂分支判断显卡则是几千人的流水线工人每个人只会做极简单的加减乘除但胜在人多、同时开工。这个差异直接决定了什么样的任务适合丢给显卡。矩阵乘法之所以是深度学习的核心就是因为它天然可以拆成海量的、彼此独立的乘加运算——正好是流水线工人最擅长的活。而像if-else密集的逻辑、递归、字符串处理这类需要“聪明人”的活显卡反而干得很差。我在实际项目里见过有人把数据预处理的一整段 Python 逻辑塞进显卡上跑结果比 CPU 慢三倍原因就在这那段逻辑里全是条件分支几千个工人挤在一起互相等还不如十几个精英各干各的。PyTorch 的 Tensor 之所以能吃到这个红利是因为它把“数据”和“在哪个设备上算”绑定在了一起。一个张量带着自己的device属性运算时框架会把整个 kernel 派发到对应设备上。所以 GPU 加速的本质不是“这段代码变快了”而是“这段代码被换了一个执行引擎”。搞清楚这个前提后面所有的坑才有解释的落脚点。1.2 决定加速上限的两个硬指标算力和带宽真正限制你能拿到多少加速比的是两个数字浮点算力和显存带宽。前者决定“单位时间能做多少次乘加”后者决定“单位时间能喂多少数据进去”。这两者经常是一起看的因为深度学习 workload 大多不是纯计算瓶颈而是“算一会儿、读一会儿数据”的混合型。举个具体的估算例子。假设你手上一张 12GB 显存级别的消费卡标称的 FP32 算力在十几到二十几 TFLOPS 这个量级显存带宽在几百 GB/s 这个量级。而你机器上的 CPU配上双通道内存带宽大概在几十 GB/s。也就是说光是把数据从显存搬到计算单元这件事显卡的通道宽度就是 CPU 内存通道的好几倍。这就是为什么哪怕是一个纯内存搬运型的操作比如大张量的逐元素加法显卡也能快出好几倍。但反过来这个估算也告诉你加速比的天花板在哪。如果你的操作是“读一个数、算一下、写回去”那它就是个带宽受限的操作加速比基本被带宽比例锁死如果是大矩阵相乘这种“读少量数据、做大量计算”的才能逼近算力比例。我一般在做性能预估时会先问自己一句这个操作是算得多还是搬得多答案不同预期就该不同。还有一个容易被忽略的约束显存容量。算力和带宽决定“多快”容量决定“能不能跑”。显存不够就是直接 OOM没有降级方案。所以真正做项目时我往往先看显存够不够再看速度。这个顺序搞反了会浪费大量时间在优化一个根本跑不起来的方案上。1.3 有些张量操作天生吃不到加速红利这是新手最容易产生的误解以为张量一旦搬到显卡上所有操作都会变快。事实是一批操作在显卡上反而更慢。第一类是极小的张量。一个形状为(4, 4)的矩阵乘法在 CPU 上可能只要一两微秒在显卡上却要付出一次 kernel 启动的固定开销这个开销本身就有几微秒量级加上主机到设备的调度延迟整体反而更慢。我做过一个粗糙的对照对 1000 个 8×8 的小矩阵做批量运算放显卡上跑的时间是 CPU 的两倍多。这类场景的正确做法是尽量把小张量合并成大张量批量处理而不是逐个丢给显卡。第二类是频繁的主机与设备之间来回搬运。每一次.cpu()和.to(cuda)都是一次跨总线的拷贝开销远高于张量本身的计算。我见过最典型的反面写法是在训练循环里对每个 batch 的每个字段单独.to(cuda)一路搬十几次改成在数据加载阶段一次性搬到显卡上整体吞吐直接涨了一截。第三类是强依赖 Python 控制流的操作。张量在显卡上算但for循环、.item()取值、if判断这些是在主机侧执行的。每次.item()都会强制同步把显卡的异步执行管线打断。这个坑我在写评估指标统计时踩过循环里每算一个 batch 就.item()一次累加到 Python 变量最后发现显卡利用率只有三成。改成在显卡上累加、最后一次性取回问题就没了。2. 把张量搬到显卡上的几条路径代价各不同2.1 四种写法看起来一样语义上差得远在 PyTorch 里把张量送上显卡至少有四种常见写法x.to(cuda)、x.cuda()、x.to(torch.device(cuda:0))、torch.tensor(data, devicecuda)。新手往往随机挑一个用实际上它们的差别不小。x.cuda()是最老派的写法硬编码了“去显卡”这件事代码里到处是它将来想切回 CPU 或者切到第二张卡就得全局替换。x.to(cuda)稍微好一点字符串形式仍然硬编码了设备名。x.to(torch.device(cuda:0))是我推荐的写法设备对象可以当参数传、可以在配置里读、可以按环境切换这是写出设备无关代码的基础。至于torch.tensor(data, device...)它是在创建时就直接在目标设备上分配省掉了一次“先在 CPU 上建、再搬过去”的过程。这个差别在小数据上无感但在创建大张量时能省下实打实的时间和一次峰值内存占用——因为“先建再搬”的瞬间CPU 和显卡上会同时存在一份完整数据。我在处理大尺寸初始化比如几亿元素的常数张量时一律用创建时就指定设备的方式。还有一点必须强调对 Tensor 调用.to()是返回新张量不是原地修改。写x.to(cuda)而不用返回值接收是新手最常见的“改了没生效”原因。而对nn.Module调用.to()则是原地修改模块内的参数这个语义差异一定要记牢。2.2 pin_memory 与非阻塞拷贝让搬运和计算重叠起来默认情况下从主机内存到显存的拷贝是同步的而且要求源内存是“可分页”的普通内存。这里有个隐藏成本显卡的 DMA 引擎无法直接访问可分页内存驱动需要先偷偷复制到一块临时的锁页缓冲区再传输。锁页内存pinned memory就是一块被锁定、不会被操作系统换出的内存DMA 可以直接读。于是就有了标准写法DataLoader(..., pin_memoryTrue)配合tensor.to(device, non_blockingTrue)。前者的作用是把 batch 数据放进锁页内存后者是告诉框架“这个拷贝可以异步进行别阻塞主机线程”。两者配合就能实现“显卡在算当前 batch 的同时下一个 batch 已经在路上”。loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) for batch in loader: x batch[x].to(device, non_blockingTrue) y batch[y].to(device, non_blockingTrue) # 前向、反向这里有个我踩过的坑non_blockingTrue只在源内存是锁页内存时才真正异步。如果忘了设pin_memoryTrue这个参数基本等于没写甚至在某些情况下还会引发难以察觉的数据一致性问题——因为主机侧可能在拷贝完成前就修改了那块内存。所以这两个参数要么一起用要么都别用只写一半是很危险的组合。另外pin_memoryTrue不是免费的午餐。锁页内存会占用系统物理内存且不可换出开太多 worker、batch 太大时机器内存会被吃掉一大块。我一般会把 worker 数控制在 CPU 核心数的一半左右batch 大小按显存反推避免两边打架。2.3 搬运时的形状、类型与连续性设备搬家的同时很多人还会顺手做类型转换比如.to(device, dtypetorch.float16)。这个用法本身没问题但要注意转换顺序带来的中间张量。.to()如果同时改变设备和精度框架内部可能先生成一个中间结果再转峰值显存会比预期高一些。还有一个更隐蔽的点内存连续性。经过transpose、permute、切片等操作之后张量的逻辑形状和物理内存布局不再一致此时它可能是“非连续”的。非连续张量搬到显卡上没问题但某些底层 kernel 对连续性有要求框架会在内部静默做一次.contiguous()拷贝。这个过程不报错、不打日志但会额外吃一份显存。我的处理习惯是在数据管线的最后、送进模型之前对确实需要连续性的张量显式调用.contiguous()并且放在搬到显卡之前做。理由很实际——在主机内存上做连续性整理出问题时更容易观察而且不会在显卡上产生瞬时双份占用。等上了显卡再整理一旦碰上显存吃紧的时刻这一下就可能成为压垮的最后一根稻草。2.4 一张表看清几种写法的适用场景写法是否原地修改适用场景主要风险x.cuda()否返回新张量临时脚本、快速验证硬编码设备难以迁移x.to(cuda)否单卡固定环境多卡下默认去 0 号卡易踩坑x.to(device_obj)否正式项目、多卡、需切换需要自己维护 device 变量torch.tensor(data, device...)不适用大张量初始化数据源仍需先存在于主机侧module.to(device)是原地改参数模型迁移之后再取参数引用可能失效最后一行值得单独说一句。module.to(device)会原地替换模块内部的参数张量对象。如果你在调用之前保存了某个参数的引用比如w model.fc.weight调用之后这个引用可能还指向旧对象。我在写自定义优化器逻辑时被这个坑过一次表现是“优化器更新了一个已经不参与前向的参数”loss 完全不动排查了很久。稳妥的做法是模型迁移设备的动作放在构建优化器之前。3. 把 device 当成一个变量来管理而不是到处硬编码3.1 一个 device 变量贯穿全流程写过几个项目之后我基本固定了一套写法在程序入口处决定设备然后把它作为参数往下传。判断逻辑也很固定——有可用显卡就用显卡没有就退回 CPU同时支持用环境变量强制指定。import os import torch def pick_device(): forced os.environ.get(FORCE_DEVICE) if forced: return torch.device(forced) if torch.cuda.is_available(): return torch.device(cuda:0) return torch.device(cpu) device pick_device()这段代码看起来平淡但它解决了一个非常实际的问题同一份代码要在开发机、训练机、别人的机器上都能跑。硬编码.cuda()的脚本一旦拿到没有可用显卡的环境上会在第一行就崩掉而且是那种带着长长堆栈的崩溃对不熟悉的人非常不友好。设备变量定下来之后所有张量创建、模型迁移、损失计算、指标统计全部围绕它走。我在 code review 时有个简单的检查标准搜索整个仓库里.cuda()出现的次数除了极少数确实需要指定具体卡号的地方其他都应该为零。3.2 多卡环境下“默认 0 号卡”的隐性坑x.to(cuda)和x.cuda()在不指定编号时都会去当前默认设备。这个默认设备可以通过torch.cuda.set_device()改变也可能被外部框架改掉。于是就有了那种“明明指定了第二张卡显存却还是溢出在第一张卡上”的诡异现象。更麻烦的是张量与模型不在同一张卡上时的报错。典型信息长这样Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!。看到这个报错第一反应不该是去改那行报错的代码而应该顺着数据流往上找哪一步用了字符串形式的cuda而没带编号。我的经验是只要项目有可能用到多卡从第一天起就全部用带编号的设备对象绝不写裸的cuda。这条规矩看着死板但省下的排查时间非常可观。尤其是当代码里混用了第三方库、每个库对默认设备的处理还不完全一致时显式编号几乎是唯一可靠的做法。3.3 加载权重时的设备错位另一个高频场景是加载保存的权重。在显卡上训练、保存然后在只有 CPU 的机器上加载如果不做处理会直接报“试图在无可用显卡的环境上反序列化显卡张量”的错误。标准解法是torch.load(path, map_locationdevice)。map_location参数会在反序列化时把张量重映射到指定设备避免“先建在显卡上再搬”的中间过程。它还可以接受一个函数做更精细的逐张量映射不过日常项目里直接传设备对象就够了。这里有个细节值得注意加载完之后权重的设备由map_location决定而模型的设备由model.to(device)决定这两件事是独立的。我见过有人只做了一半结果load_state_dict报错说键对不上——其实不是键对不上是两边设备不一致导致的加载失败被包装成了别的错误信息。稳妥顺序是先建模型、先model.to(device)、再加载权重。另外还有一个便利写法model.load_state_dict(state)之后如果想验证设备是否一致可以取任意一个参数看它的.device比逐个检查省事得多。我在每个训练脚本的启动阶段都会打一行日志把模型参数的设备、dtype 和数量打出来出问题时这行日志往往能直接定位。4. 显存这本账先会算再会查4.1 一个张量到底吃多少显存显存占用的算法很朴素元素个数 × 每个元素的字节数。float32 是 4 字节float16 和 bfloat16 都是 2 字节int64 是 8 字节int32 是 4 字节bool 是 1 字节。一个形状为(1024, 1024)的 float32 张量元素数约 104 万占用约 4MB。看着不多但训练时的显存大头从来不是单个张量。真正的账要这么算。假设模型有 1000 万个参数用 float32 训练项目计算方式占用模型参数1e7 × 4B约 40MB梯度与参数同尺寸约 40MB优化器一阶动量与参数同尺寸约 40MB优化器二阶动量与参数同尺寸约 40MB中间激活值与 batch、序列长度、层数强相关常常是最大项也就是说光是参数加优化器状态就已经是参数本身的四倍。而真正的显存杀手是激活值——它跟 batch size 和输入长度是线性关系而且在前向过程中要保留下来给反向用。这就解释了为什么“模型参数量不大但就是 OOM”问题往往出在 batch 或输入尺寸上而不是模型本身。我一般在动手调参之前先按上表粗算一遍得出“理论下限”然后用实际运行时的占用去对比。两者差距过大说明有额外的浪费差距过小说明显存已经绷得很紧该考虑梯度累积或者降低 batch 了。4.2 allocated 和 reserved 是两回事PyTorch 的缓存分配器有两个关键指标torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。前者是当前真正被张量占用的字节数后者是分配器向驱动申请下来、留在自己池子里的总字节数。这个区别非常重要。我第一次看到reserved远大于allocated时以为是显存泄漏折腾了半天各种del最后才明白张量被释放后分配器并不会立刻把内存还给驱动而是留在池子里等下次复用。这是性能优化——避免频繁向驱动申请和归还。所以reserved高但allocated低是完全正常的状态不需要处理。真正需要警惕的是allocated随迭代单调上升。这种情况通常意味着有张量被意外持有比如在列表里持续append中间结果而没有清理在计算图里保留了带梯度的张量导致整条计算图无法释放典型如把 loss 累加进一个列表闭包或全局变量意外引用了大张量。排查时我会在训练循环里每隔若干步打一次allocated和reserved画个简单曲线上升趋势一眼就能看出来。这个方法比任何静态分析都管用。4.3 empty_cache 不是万能药用错反而更慢torch.cuda.empty_cache()的作用是把分配器池子里空闲的缓存还给驱动。注意“空闲”两个字它不会释放任何正在被张量使用的显存。所以那种“OOM 了就调一下 empty_cache”的做法绝大多数情况下毫无效果。更要命的是性能代价。释放缓存之后下一次需要显存时又得重新向驱动申请这个往返有成本。如果在训练循环里每步都调一次整体速度会明显下降。我在一个项目里见过有人在每个 batch 末尾都调训练时间凭空多了两成。我的使用原则很窄只在“加载大模型之前”“切换任务之前”“确实观察到 reserved 长期远大于 allocated 且即将触发 OOM”这三种情况下调用。日常训练循环里一次都不该出现。顺便说一下碎片问题。长时间训练、反复分配不同尺寸的张量会让缓存池出现大量小空洞总空闲显存够但找不到连续的大块于是 OOM。缓解手段包括尽量固定 batch 形状避免动态尺寸、把不同尺寸的张量分配集中在训练开始前完成、以及在极端情况下重启进程。真正做长周期训练时我会把输入尺寸固定住就是为了规避碎片。4.4 一次真实的显存增长排查链路说个具体的。有次训练跑了大概两千步之后 OOM前两千步一切正常显存曲线缓慢上升。第一步我打印了allocated和reserved的差值发现 reserved 基本平稳allocated 在涨说明确实是张量没释放不是碎片。第二步检查训练循环里所有append和字典写入。发现指标统计那里写成了losses.append(loss)——而loss是带计算图的张量。这个列表持有 lossloss 持有整条反向图于是每一步的计算图都没被释放。改成losses.append(loss.item())之后曲线立刻变平。第三步顺手做了个优化.item()每次调用都会触发一次设备同步在循环里频繁调用同样伤性能。所以最终改成在显卡上累加、循环结束统一取回total torch.zeros((), devicedevice) for batch in loader: loss compute_loss(batch) total loss.detach() avg (total / len(loader)).item()detach()是关键它把张量从计算图上摘下来不再持有反向所需的中间结果。这个写法既避免了显存泄漏又避免了每步同步。5. 精度与设备的组合是报错最密集的地带5.1 半精度的收益和边界把张量从 float32 换成 float16 或 bfloat16最直接的好处是显存占用减半、带宽需求减半。在带宽受限的操作上这几乎等于免费提速。但半精度不是随便一转就完事。float16 的动态范围窄容易出现溢出或下溢bfloat16 动态范围和 float32 一致但尾数少精度略低但在训练中通常够用。我自己的选择习惯是推理和存储用 float16训练优先考虑 bfloat16硬件支持的前提下。这个取舍没有绝对答案取决于具体任务对数值稳定性的敏感程度。更重要的是不能简单地把整个模型.half()了事。归一化层、损失计算、某些归约操作在半精度下容易出问题。稳妥的做法是用自动混合精度让框架自己决定每个算子用什么精度而不是手动一刀切。scaler torch.cuda.amp.GradScaler() for batch in loader: optimizer.zero_grad(set_to_noneTrue) with torch.cuda.amp.autocast(dtypetorch.float16): out model(batch[x].to(device, non_blockingTrue)) loss criterion(out, batch[y].to(device, non_blockingTrue)) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是给损失乘一个缩放因子避免梯度在半精度下下溢成零更新参数前再除回来。这段代码我基本是原样复用的唯一会调的是初始缩放因子——如果训练早期频繁出现“跳过更新”的警告说明缩放因子在反复调整可以考虑手动指定一个更合适的初值。5.2 那些“只在显卡上出现”的类型问题有一类报错很迷惑同样的代码在 CPU 上跑得好好的搬到显卡上就报类型不支持。常见原因有两个。第一个是设备本身不支持某个 dtype。某些老一些的硬件对 bfloat16 的支持不完整部分算子会退回或直接报错。这种情况下最省事的做法是先查一下当前设备的计算能力再决定用哪种精度而不是写死。第二个是索引张量的类型。用张量做索引时比如x[idx]idx必须是整型。如果它是 floatCPU 上某些路径可能容忍显卡上则直接报错。我在做数据增强时被这个坑过一个坐标计算出来是浮点忘了转long就去索引CPU 上恰好没触发问题显卡上每一步都崩。第三个是设备不匹配。前面提过报错信息这里补一个排查技巧在报错位置的前一行打印所有相关张量的.device和.dtype。听起来笨但比逐行推理快得多。我现在的习惯是在模型forward开头加一段只在调试模式下启用的断言if DEBUG: for name, t in {x: x, mask: mask}.items(): assert t.device self.device, f{name} on {t.device}这几行代码上线时可以整体去掉开发阶段能省掉大量来回。5.3 常见报错与定位方向对照报错关键词最可能的原因优先检查found at least two devices张量分散在不同设备是否有裸cuda未带编号Torch not compiled with CUDA enabled安装的是纯 CPU 版本环境里框架的构建版本CUDA out of memory显存不足batch 大小、激活值、是否有泄漏expected scalar type Half but found Float混精度的算子边界归一化层、损失函数的输入 dtypeindex tensor must have dtype long索引用浮点索引张量的.long()转换这张表是我自己攒的遇到报错先对一眼能省掉不少搜索时间。需要说明的是同一个报错可能有多种成因表格给的是“最可能”不是“一定”。真正定位还得靠打印和二分法。6. 一段能复现的基准测试以及加速比被稀释的真相6.1 基准测试的三个必备动作写性能对比代码时有三个动作缺一不可缺了任何一个测出来的数字都不可信。第一个是预热。显卡第一次执行某个 kernel 时要做上下文初始化、编译、加载第一次的耗时往往是稳定状态的数倍。不预热直接计时测到的是冷启动成本。第二个是同步。显卡上的操作默认是异步派发的主机侧发起调用后立刻返回实际计算还在排队。如果不调torch.cuda.synchronize()就在主机侧读计时器测到的只是“派发完成”的时间不是“计算完成”的时间。第三个是多次取中位数。单次测量受各种噪声影响取多次的中位数比取平均更稳平均值容易被个别异常值拉偏。import time import torch def bench(fn, warmup10, repeat50, devicecuda): for _ in range(warmup): fn() if device.startswith(cuda): torch.cuda.synchronize() start_evt torch.cuda.Event(enable_timingTrue) end_evt torch.cuda.Event(enable_timingTrue) times [] for _ in range(repeat): start_evt.record() fn() end_evt.record() torch.cuda.synchronize() times.append(start_evt.elapsed_time(end_evt)) else: times [] for _ in range(repeat): t0 time.perf_counter() fn() times.append((time.perf_counter() - t0) * 1000) times.sort() return times[len(times) // 2]用事件计时比用time.perf_counter更准因为它测的是显卡上的真实执行时间不受主机侧调度抖动影响。这段代码我建议直接收进工具箱后面所有性能对比都用它。6.2 小张量上显卡真的会输拿矩阵乘法做对照尺寸从 32 一路加到 4096大致会看到这样的趋势矩阵尺寸CPU 相对耗时显卡相对耗时观察32 × 321.03~5显卡明显更慢128 × 1281.00.5 左右大致打平512 × 5121.00.1 左右显卡开始领先2048 × 20481.00.02 以下显卡碾压表中的倍数是相对值具体数字随硬件差异很大但趋势是稳定的。这个表解释了我开头那个“改到显卡上反而更慢”的经历。当时我的模型很小、batch 也很小绝大部分时间花在 kernel 启动和跨设备搬运上真正的计算量根本不值一提。这种情况下最优解可能压根不是显卡而是把 CPU 上的并行度调好或者干脆把小张量批量化。6.3 端到端训练里的加速比为什么总比基准测试低很多人测完矩阵乘法看到几十倍加速兴冲冲去跑训练结果端到端只快了三四倍。这不是显卡不行而是被几个环节稀释了。稀释点一数据加载。如果数据预处理在 CPU 上做显卡算得再快也得等。这时候要看的不是显卡利用率而是数据加载的吞吐是否跟得上。开多 worker、把解码和增强做得更轻、或者把部分变换也搬到显卡上做比如用显卡上的张量操作做归一化都是常见手段。稀释点二主机与设备的往返。每个 batch 搬几次、有没有不必要的.cpu()取回、训练循环里有没有频繁.item()这些都在偷偷吃掉加速收益。稀释点三同步点。日志打印、指标记录、学习率调度里如果有隐式或显式的同步操作异步管线就被打断显卡会有空转时间。稀释点四小算子过多。一层里几十个小算子每个都要一次 kernel 启动累加起来开销可观。这种情况下算子融合的收益比换硬件更直接。我一般的排查顺序是先看显卡利用率曲线如果长期低于五成基本可以断定瓶颈在数据侧如果利用率高但加速比仍低那就是算法本身的计算密度不够属于“无解”只能改方案。6.4 两个容易被忽略的开关最后补两个实践中的小开关。torch.backends.cudnn.benchmark True会让框架在固定输入尺寸时自动搜一遍最快的卷积实现代价是第一次会慢一点。输入尺寸固定、训练步数多的场景值得打开输入尺寸每步都变的场景别开搜索成本会把收益吃光。另一个是矩阵乘法的精度模式。较新的硬件上默认会启用一种折中的浮点模式用少量精度换显著速度。如果你的任务对数值非常敏感可以显式关掉做对照实验如果只是训练一个普通模型保持默认通常没问题。这两个开关我都建议在正式训练前花十分钟做 A/B 对照用数据决定开还是关而不是凭感觉。我个人在几个项目里反复验证下来最值得投入时间的其实不是这些开关而是前面那几件事把设备管理收敛到一个变量、把数据搬运的路径理顺、把显存占用的账算清楚。这三件事做到位加速比自然就上来了这三件事没做好再怎么调开关都是在漏水的桶里加水。至于具体的数字不同硬件、不同模型差异太大我更愿意在自己的机器上老老实实跑一遍基准而不是照抄别人的结论——毕竟上一次照抄我就多花了一个下午。 
长河君 · 2026/10/27 2:33:19 · 阅读 1.2k PyTorch Tensor GPU加速：设备管理、显存优化与基准测试
我第一次把训练脚本从 CPU 挪到显卡上跑是很多年前的事了。当时满心以为改一行.cuda()就能看到几十倍的提速结果跑完一个 epoch计时器上的数字比 CPU 版本还难看。后来花了整整一个下午才发现问题根本不在显卡而在我根本没搞清楚 PyTorch 里 Tensor张量和 GPU 加速之间的关系哪些操作真的被搬到显卡上执行了哪些数据还老老实实躺在内存里被反复搬运哪些地方因为一个小小的设备不匹配在偷偷做隐式拷贝。这件事之后我养成了一个习惯——凡是涉及device的代码一律当成一等公民来对待而不是随手加个.to(cuda)了事。这篇内容聊的就是 PyTorch 数据结构里最基础也最容易被轻视的一环Tensor 的 GPU 加速。它不是什么高深话题但真正把它用明白的人并不多。适合刚接触 PyTorch、想把模型跑到显卡上的新手也适合已经能跑通训练、但总是被显存溢出和设备报错折腾的中级使用者。我会从“为什么快”讲到“怎么搬”再讲到“显存这本账怎么算”和“精度怎么配”最后给一段可复现的基准测试写法。全程都是我自己踩过坑之后总结的东西能直接抄去用。1. 先把“为什么快”这笔账算清楚1.1 显卡不是“更快的 CPU”而是“宽得多的 CPU”很多人对 GPU 加速的第一印象是“显卡主频高”。恰恰相反一张中端显卡的单核频率往往还不如同时代的桌面 CPU。它快的地方在别处CPU 可能只有十几个核心而显卡上有几千个计算单元。你可以把 CPU 想象成一个十几个人的精英小队每个人都能独当一面、能处理复杂分支判断显卡则是几千人的流水线工人每个人只会做极简单的加减乘除但胜在人多、同时开工。这个差异直接决定了什么样的任务适合丢给显卡。矩阵乘法之所以是深度学习的核心就是因为它天然可以拆成海量的、彼此独立的乘加运算——正好是流水线工人最擅长的活。而像if-else密集的逻辑、递归、字符串处理这类需要“聪明人”的活显卡反而干得很差。我在实际项目里见过有人把数据预处理的一整段 Python 逻辑塞进显卡上跑结果比 CPU 慢三倍原因就在这那段逻辑里全是条件分支几千个工人挤在一起互相等还不如十几个精英各干各的。PyTorch 的 Tensor 之所以能吃到这个红利是因为它把“数据”和“在哪个设备上算”绑定在了一起。一个张量带着自己的device属性运算时框架会把整个 kernel 派发到对应设备上。所以 GPU 加速的本质不是“这段代码变快了”而是“这段代码被换了一个执行引擎”。搞清楚这个前提后面所有的坑才有解释的落脚点。1.2 决定加速上限的两个硬指标算力和带宽真正限制你能拿到多少加速比的是两个数字浮点算力和显存带宽。前者决定“单位时间能做多少次乘加”后者决定“单位时间能喂多少数据进去”。这两者经常是一起看的因为深度学习 workload 大多不是纯计算瓶颈而是“算一会儿、读一会儿数据”的混合型。举个具体的估算例子。假设你手上一张 12GB 显存级别的消费卡标称的 FP32 算力在十几到二十几 TFLOPS 这个量级显存带宽在几百 GB/s 这个量级。而你机器上的 CPU配上双通道内存带宽大概在几十 GB/s。也就是说光是把数据从显存搬到计算单元这件事显卡的通道宽度就是 CPU 内存通道的好几倍。这就是为什么哪怕是一个纯内存搬运型的操作比如大张量的逐元素加法显卡也能快出好几倍。但反过来这个估算也告诉你加速比的天花板在哪。如果你的操作是“读一个数、算一下、写回去”那它就是个带宽受限的操作加速比基本被带宽比例锁死如果是大矩阵相乘这种“读少量数据、做大量计算”的才能逼近算力比例。我一般在做性能预估时会先问自己一句这个操作是算得多还是搬得多答案不同预期就该不同。还有一个容易被忽略的约束显存容量。算力和带宽决定“多快”容量决定“能不能跑”。显存不够就是直接 OOM没有降级方案。所以真正做项目时我往往先看显存够不够再看速度。这个顺序搞反了会浪费大量时间在优化一个根本跑不起来的方案上。1.3 有些张量操作天生吃不到加速红利这是新手最容易产生的误解以为张量一旦搬到显卡上所有操作都会变快。事实是一批操作在显卡上反而更慢。第一类是极小的张量。一个形状为(4, 4)的矩阵乘法在 CPU 上可能只要一两微秒在显卡上却要付出一次 kernel 启动的固定开销这个开销本身就有几微秒量级加上主机到设备的调度延迟整体反而更慢。我做过一个粗糙的对照对 1000 个 8×8 的小矩阵做批量运算放显卡上跑的时间是 CPU 的两倍多。这类场景的正确做法是尽量把小张量合并成大张量批量处理而不是逐个丢给显卡。第二类是频繁的主机与设备之间来回搬运。每一次.cpu()和.to(cuda)都是一次跨总线的拷贝开销远高于张量本身的计算。我见过最典型的反面写法是在训练循环里对每个 batch 的每个字段单独.to(cuda)一路搬十几次改成在数据加载阶段一次性搬到显卡上整体吞吐直接涨了一截。第三类是强依赖 Python 控制流的操作。张量在显卡上算但for循环、.item()取值、if判断这些是在主机侧执行的。每次.item()都会强制同步把显卡的异步执行管线打断。这个坑我在写评估指标统计时踩过循环里每算一个 batch 就.item()一次累加到 Python 变量最后发现显卡利用率只有三成。改成在显卡上累加、最后一次性取回问题就没了。2. 把张量搬到显卡上的几条路径代价各不同2.1 四种写法看起来一样语义上差得远在 PyTorch 里把张量送上显卡至少有四种常见写法x.to(cuda)、x.cuda()、x.to(torch.device(cuda:0))、torch.tensor(data, devicecuda)。新手往往随机挑一个用实际上它们的差别不小。x.cuda()是最老派的写法硬编码了“去显卡”这件事代码里到处是它将来想切回 CPU 或者切到第二张卡就得全局替换。x.to(cuda)稍微好一点字符串形式仍然硬编码了设备名。x.to(torch.device(cuda:0))是我推荐的写法设备对象可以当参数传、可以在配置里读、可以按环境切换这是写出设备无关代码的基础。至于torch.tensor(data, device...)它是在创建时就直接在目标设备上分配省掉了一次“先在 CPU 上建、再搬过去”的过程。这个差别在小数据上无感但在创建大张量时能省下实打实的时间和一次峰值内存占用——因为“先建再搬”的瞬间CPU 和显卡上会同时存在一份完整数据。我在处理大尺寸初始化比如几亿元素的常数张量时一律用创建时就指定设备的方式。还有一点必须强调对 Tensor 调用.to()是返回新张量不是原地修改。写x.to(cuda)而不用返回值接收是新手最常见的“改了没生效”原因。而对nn.Module调用.to()则是原地修改模块内的参数这个语义差异一定要记牢。2.2 pin_memory 与非阻塞拷贝让搬运和计算重叠起来默认情况下从主机内存到显存的拷贝是同步的而且要求源内存是“可分页”的普通内存。这里有个隐藏成本显卡的 DMA 引擎无法直接访问可分页内存驱动需要先偷偷复制到一块临时的锁页缓冲区再传输。锁页内存pinned memory就是一块被锁定、不会被操作系统换出的内存DMA 可以直接读。于是就有了标准写法DataLoader(..., pin_memoryTrue)配合tensor.to(device, non_blockingTrue)。前者的作用是把 batch 数据放进锁页内存后者是告诉框架“这个拷贝可以异步进行别阻塞主机线程”。两者配合就能实现“显卡在算当前 batch 的同时下一个 batch 已经在路上”。loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) for batch in loader: x batch[x].to(device, non_blockingTrue) y batch[y].to(device, non_blockingTrue) # 前向、反向这里有个我踩过的坑non_blockingTrue只在源内存是锁页内存时才真正异步。如果忘了设pin_memoryTrue这个参数基本等于没写甚至在某些情况下还会引发难以察觉的数据一致性问题——因为主机侧可能在拷贝完成前就修改了那块内存。所以这两个参数要么一起用要么都别用只写一半是很危险的组合。另外pin_memoryTrue不是免费的午餐。锁页内存会占用系统物理内存且不可换出开太多 worker、batch 太大时机器内存会被吃掉一大块。我一般会把 worker 数控制在 CPU 核心数的一半左右batch 大小按显存反推避免两边打架。2.3 搬运时的形状、类型与连续性设备搬家的同时很多人还会顺手做类型转换比如.to(device, dtypetorch.float16)。这个用法本身没问题但要注意转换顺序带来的中间张量。.to()如果同时改变设备和精度框架内部可能先生成一个中间结果再转峰值显存会比预期高一些。还有一个更隐蔽的点内存连续性。经过transpose、permute、切片等操作之后张量的逻辑形状和物理内存布局不再一致此时它可能是“非连续”的。非连续张量搬到显卡上没问题但某些底层 kernel 对连续性有要求框架会在内部静默做一次.contiguous()拷贝。这个过程不报错、不打日志但会额外吃一份显存。我的处理习惯是在数据管线的最后、送进模型之前对确实需要连续性的张量显式调用.contiguous()并且放在搬到显卡之前做。理由很实际——在主机内存上做连续性整理出问题时更容易观察而且不会在显卡上产生瞬时双份占用。等上了显卡再整理一旦碰上显存吃紧的时刻这一下就可能成为压垮的最后一根稻草。2.4 一张表看清几种写法的适用场景写法是否原地修改适用场景主要风险x.cuda()否返回新张量临时脚本、快速验证硬编码设备难以迁移x.to(cuda)否单卡固定环境多卡下默认去 0 号卡易踩坑x.to(device_obj)否正式项目、多卡、需切换需要自己维护 device 变量torch.tensor(data, device...)不适用大张量初始化数据源仍需先存在于主机侧module.to(device)是原地改参数模型迁移之后再取参数引用可能失效最后一行值得单独说一句。module.to(device)会原地替换模块内部的参数张量对象。如果你在调用之前保存了某个参数的引用比如w model.fc.weight调用之后这个引用可能还指向旧对象。我在写自定义优化器逻辑时被这个坑过一次表现是“优化器更新了一个已经不参与前向的参数”loss 完全不动排查了很久。稳妥的做法是模型迁移设备的动作放在构建优化器之前。3. 把 device 当成一个变量来管理而不是到处硬编码3.1 一个 device 变量贯穿全流程写过几个项目之后我基本固定了一套写法在程序入口处决定设备然后把它作为参数往下传。判断逻辑也很固定——有可用显卡就用显卡没有就退回 CPU同时支持用环境变量强制指定。import os import torch def pick_device(): forced os.environ.get(FORCE_DEVICE) if forced: return torch.device(forced) if torch.cuda.is_available(): return torch.device(cuda:0) return torch.device(cpu) device pick_device()这段代码看起来平淡但它解决了一个非常实际的问题同一份代码要在开发机、训练机、别人的机器上都能跑。硬编码.cuda()的脚本一旦拿到没有可用显卡的环境上会在第一行就崩掉而且是那种带着长长堆栈的崩溃对不熟悉的人非常不友好。设备变量定下来之后所有张量创建、模型迁移、损失计算、指标统计全部围绕它走。我在 code review 时有个简单的检查标准搜索整个仓库里.cuda()出现的次数除了极少数确实需要指定具体卡号的地方其他都应该为零。3.2 多卡环境下“默认 0 号卡”的隐性坑x.to(cuda)和x.cuda()在不指定编号时都会去当前默认设备。这个默认设备可以通过torch.cuda.set_device()改变也可能被外部框架改掉。于是就有了那种“明明指定了第二张卡显存却还是溢出在第一张卡上”的诡异现象。更麻烦的是张量与模型不在同一张卡上时的报错。典型信息长这样Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!。看到这个报错第一反应不该是去改那行报错的代码而应该顺着数据流往上找哪一步用了字符串形式的cuda而没带编号。我的经验是只要项目有可能用到多卡从第一天起就全部用带编号的设备对象绝不写裸的cuda。这条规矩看着死板但省下的排查时间非常可观。尤其是当代码里混用了第三方库、每个库对默认设备的处理还不完全一致时显式编号几乎是唯一可靠的做法。3.3 加载权重时的设备错位另一个高频场景是加载保存的权重。在显卡上训练、保存然后在只有 CPU 的机器上加载如果不做处理会直接报“试图在无可用显卡的环境上反序列化显卡张量”的错误。标准解法是torch.load(path, map_locationdevice)。map_location参数会在反序列化时把张量重映射到指定设备避免“先建在显卡上再搬”的中间过程。它还可以接受一个函数做更精细的逐张量映射不过日常项目里直接传设备对象就够了。这里有个细节值得注意加载完之后权重的设备由map_location决定而模型的设备由model.to(device)决定这两件事是独立的。我见过有人只做了一半结果load_state_dict报错说键对不上——其实不是键对不上是两边设备不一致导致的加载失败被包装成了别的错误信息。稳妥顺序是先建模型、先model.to(device)、再加载权重。另外还有一个便利写法model.load_state_dict(state)之后如果想验证设备是否一致可以取任意一个参数看它的.device比逐个检查省事得多。我在每个训练脚本的启动阶段都会打一行日志把模型参数的设备、dtype 和数量打出来出问题时这行日志往往能直接定位。4. 显存这本账先会算再会查4.1 一个张量到底吃多少显存显存占用的算法很朴素元素个数 × 每个元素的字节数。float32 是 4 字节float16 和 bfloat16 都是 2 字节int64 是 8 字节int32 是 4 字节bool 是 1 字节。一个形状为(1024, 1024)的 float32 张量元素数约 104 万占用约 4MB。看着不多但训练时的显存大头从来不是单个张量。真正的账要这么算。假设模型有 1000 万个参数用 float32 训练项目计算方式占用模型参数1e7 × 4B约 40MB梯度与参数同尺寸约 40MB优化器一阶动量与参数同尺寸约 40MB优化器二阶动量与参数同尺寸约 40MB中间激活值与 batch、序列长度、层数强相关常常是最大项也就是说光是参数加优化器状态就已经是参数本身的四倍。而真正的显存杀手是激活值——它跟 batch size 和输入长度是线性关系而且在前向过程中要保留下来给反向用。这就解释了为什么“模型参数量不大但就是 OOM”问题往往出在 batch 或输入尺寸上而不是模型本身。我一般在动手调参之前先按上表粗算一遍得出“理论下限”然后用实际运行时的占用去对比。两者差距过大说明有额外的浪费差距过小说明显存已经绷得很紧该考虑梯度累积或者降低 batch 了。4.2 allocated 和 reserved 是两回事PyTorch 的缓存分配器有两个关键指标torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。前者是当前真正被张量占用的字节数后者是分配器向驱动申请下来、留在自己池子里的总字节数。这个区别非常重要。我第一次看到reserved远大于allocated时以为是显存泄漏折腾了半天各种del最后才明白张量被释放后分配器并不会立刻把内存还给驱动而是留在池子里等下次复用。这是性能优化——避免频繁向驱动申请和归还。所以reserved高但allocated低是完全正常的状态不需要处理。真正需要警惕的是allocated随迭代单调上升。这种情况通常意味着有张量被意外持有比如在列表里持续append中间结果而没有清理在计算图里保留了带梯度的张量导致整条计算图无法释放典型如把 loss 累加进一个列表闭包或全局变量意外引用了大张量。排查时我会在训练循环里每隔若干步打一次allocated和reserved画个简单曲线上升趋势一眼就能看出来。这个方法比任何静态分析都管用。4.3 empty_cache 不是万能药用错反而更慢torch.cuda.empty_cache()的作用是把分配器池子里空闲的缓存还给驱动。注意“空闲”两个字它不会释放任何正在被张量使用的显存。所以那种“OOM 了就调一下 empty_cache”的做法绝大多数情况下毫无效果。更要命的是性能代价。释放缓存之后下一次需要显存时又得重新向驱动申请这个往返有成本。如果在训练循环里每步都调一次整体速度会明显下降。我在一个项目里见过有人在每个 batch 末尾都调训练时间凭空多了两成。我的使用原则很窄只在“加载大模型之前”“切换任务之前”“确实观察到 reserved 长期远大于 allocated 且即将触发 OOM”这三种情况下调用。日常训练循环里一次都不该出现。顺便说一下碎片问题。长时间训练、反复分配不同尺寸的张量会让缓存池出现大量小空洞总空闲显存够但找不到连续的大块于是 OOM。缓解手段包括尽量固定 batch 形状避免动态尺寸、把不同尺寸的张量分配集中在训练开始前完成、以及在极端情况下重启进程。真正做长周期训练时我会把输入尺寸固定住就是为了规避碎片。4.4 一次真实的显存增长排查链路说个具体的。有次训练跑了大概两千步之后 OOM前两千步一切正常显存曲线缓慢上升。第一步我打印了allocated和reserved的差值发现 reserved 基本平稳allocated 在涨说明确实是张量没释放不是碎片。第二步检查训练循环里所有append和字典写入。发现指标统计那里写成了losses.append(loss)——而loss是带计算图的张量。这个列表持有 lossloss 持有整条反向图于是每一步的计算图都没被释放。改成losses.append(loss.item())之后曲线立刻变平。第三步顺手做了个优化.item()每次调用都会触发一次设备同步在循环里频繁调用同样伤性能。所以最终改成在显卡上累加、循环结束统一取回total torch.zeros((), devicedevice) for batch in loader: loss compute_loss(batch) total loss.detach() avg (total / len(loader)).item()detach()是关键它把张量从计算图上摘下来不再持有反向所需的中间结果。这个写法既避免了显存泄漏又避免了每步同步。5. 精度与设备的组合是报错最密集的地带5.1 半精度的收益和边界把张量从 float32 换成 float16 或 bfloat16最直接的好处是显存占用减半、带宽需求减半。在带宽受限的操作上这几乎等于免费提速。但半精度不是随便一转就完事。float16 的动态范围窄容易出现溢出或下溢bfloat16 动态范围和 float32 一致但尾数少精度略低但在训练中通常够用。我自己的选择习惯是推理和存储用 float16训练优先考虑 bfloat16硬件支持的前提下。这个取舍没有绝对答案取决于具体任务对数值稳定性的敏感程度。更重要的是不能简单地把整个模型.half()了事。归一化层、损失计算、某些归约操作在半精度下容易出问题。稳妥的做法是用自动混合精度让框架自己决定每个算子用什么精度而不是手动一刀切。scaler torch.cuda.amp.GradScaler() for batch in loader: optimizer.zero_grad(set_to_noneTrue) with torch.cuda.amp.autocast(dtypetorch.float16): out model(batch[x].to(device, non_blockingTrue)) loss criterion(out, batch[y].to(device, non_blockingTrue)) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是给损失乘一个缩放因子避免梯度在半精度下下溢成零更新参数前再除回来。这段代码我基本是原样复用的唯一会调的是初始缩放因子——如果训练早期频繁出现“跳过更新”的警告说明缩放因子在反复调整可以考虑手动指定一个更合适的初值。5.2 那些“只在显卡上出现”的类型问题有一类报错很迷惑同样的代码在 CPU 上跑得好好的搬到显卡上就报类型不支持。常见原因有两个。第一个是设备本身不支持某个 dtype。某些老一些的硬件对 bfloat16 的支持不完整部分算子会退回或直接报错。这种情况下最省事的做法是先查一下当前设备的计算能力再决定用哪种精度而不是写死。第二个是索引张量的类型。用张量做索引时比如x[idx]idx必须是整型。如果它是 floatCPU 上某些路径可能容忍显卡上则直接报错。我在做数据增强时被这个坑过一个坐标计算出来是浮点忘了转long就去索引CPU 上恰好没触发问题显卡上每一步都崩。第三个是设备不匹配。前面提过报错信息这里补一个排查技巧在报错位置的前一行打印所有相关张量的.device和.dtype。听起来笨但比逐行推理快得多。我现在的习惯是在模型forward开头加一段只在调试模式下启用的断言if DEBUG: for name, t in {x: x, mask: mask}.items(): assert t.device self.device, f{name} on {t.device}这几行代码上线时可以整体去掉开发阶段能省掉大量来回。5.3 常见报错与定位方向对照报错关键词最可能的原因优先检查found at least two devices张量分散在不同设备是否有裸cuda未带编号Torch not compiled with CUDA enabled安装的是纯 CPU 版本环境里框架的构建版本CUDA out of memory显存不足batch 大小、激活值、是否有泄漏expected scalar type Half but found Float混精度的算子边界归一化层、损失函数的输入 dtypeindex tensor must have dtype long索引用浮点索引张量的.long()转换这张表是我自己攒的遇到报错先对一眼能省掉不少搜索时间。需要说明的是同一个报错可能有多种成因表格给的是“最可能”不是“一定”。真正定位还得靠打印和二分法。6. 一段能复现的基准测试以及加速比被稀释的真相6.1 基准测试的三个必备动作写性能对比代码时有三个动作缺一不可缺了任何一个测出来的数字都不可信。第一个是预热。显卡第一次执行某个 kernel 时要做上下文初始化、编译、加载第一次的耗时往往是稳定状态的数倍。不预热直接计时测到的是冷启动成本。第二个是同步。显卡上的操作默认是异步派发的主机侧发起调用后立刻返回实际计算还在排队。如果不调torch.cuda.synchronize()就在主机侧读计时器测到的只是“派发完成”的时间不是“计算完成”的时间。第三个是多次取中位数。单次测量受各种噪声影响取多次的中位数比取平均更稳平均值容易被个别异常值拉偏。import time import torch def bench(fn, warmup10, repeat50, devicecuda): for _ in range(warmup): fn() if device.startswith(cuda): torch.cuda.synchronize() start_evt torch.cuda.Event(enable_timingTrue) end_evt torch.cuda.Event(enable_timingTrue) times [] for _ in range(repeat): start_evt.record() fn() end_evt.record() torch.cuda.synchronize() times.append(start_evt.elapsed_time(end_evt)) else: times [] for _ in range(repeat): t0 time.perf_counter() fn() times.append((time.perf_counter() - t0) * 1000) times.sort() return times[len(times) // 2]用事件计时比用time.perf_counter更准因为它测的是显卡上的真实执行时间不受主机侧调度抖动影响。这段代码我建议直接收进工具箱后面所有性能对比都用它。6.2 小张量上显卡真的会输拿矩阵乘法做对照尺寸从 32 一路加到 4096大致会看到这样的趋势矩阵尺寸CPU 相对耗时显卡相对耗时观察32 × 321.03~5显卡明显更慢128 × 1281.00.5 左右大致打平512 × 5121.00.1 左右显卡开始领先2048 × 20481.00.02 以下显卡碾压表中的倍数是相对值具体数字随硬件差异很大但趋势是稳定的。这个表解释了我开头那个“改到显卡上反而更慢”的经历。当时我的模型很小、batch 也很小绝大部分时间花在 kernel 启动和跨设备搬运上真正的计算量根本不值一提。这种情况下最优解可能压根不是显卡而是把 CPU 上的并行度调好或者干脆把小张量批量化。6.3 端到端训练里的加速比为什么总比基准测试低很多人测完矩阵乘法看到几十倍加速兴冲冲去跑训练结果端到端只快了三四倍。这不是显卡不行而是被几个环节稀释了。稀释点一数据加载。如果数据预处理在 CPU 上做显卡算得再快也得等。这时候要看的不是显卡利用率而是数据加载的吞吐是否跟得上。开多 worker、把解码和增强做得更轻、或者把部分变换也搬到显卡上做比如用显卡上的张量操作做归一化都是常见手段。稀释点二主机与设备的往返。每个 batch 搬几次、有没有不必要的.cpu()取回、训练循环里有没有频繁.item()这些都在偷偷吃掉加速收益。稀释点三同步点。日志打印、指标记录、学习率调度里如果有隐式或显式的同步操作异步管线就被打断显卡会有空转时间。稀释点四小算子过多。一层里几十个小算子每个都要一次 kernel 启动累加起来开销可观。这种情况下算子融合的收益比换硬件更直接。我一般的排查顺序是先看显卡利用率曲线如果长期低于五成基本可以断定瓶颈在数据侧如果利用率高但加速比仍低那就是算法本身的计算密度不够属于“无解”只能改方案。6.4 两个容易被忽略的开关最后补两个实践中的小开关。torch.backends.cudnn.benchmark True会让框架在固定输入尺寸时自动搜一遍最快的卷积实现代价是第一次会慢一点。输入尺寸固定、训练步数多的场景值得打开输入尺寸每步都变的场景别开搜索成本会把收益吃光。另一个是矩阵乘法的精度模式。较新的硬件上默认会启用一种折中的浮点模式用少量精度换显著速度。如果你的任务对数值非常敏感可以显式关掉做对照实验如果只是训练一个普通模型保持默认通常没问题。这两个开关我都建议在正式训练前花十分钟做 A/B 对照用数据决定开还是关而不是凭感觉。我个人在几个项目里反复验证下来最值得投入时间的其实不是这些开关而是前面那几件事把设备管理收敛到一个变量、把数据搬运的路径理顺、把显存占用的账算清楚。这三件事做到位加速比自然就上来了这三件事没做好再怎么调开关都是在漏水的桶里加水。至于具体的数字不同硬件、不同模型差异太大我更愿意在自己的机器上老老实实跑一遍基准而不是照抄别人的结论——毕竟上一次照抄我就多花了一个下午。