十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PyTorch性能调优实战:从Profiling到torch.compile与分布式扩展

PyTorch性能调优实战:从Profiling到torch.compile与分布式扩展 干我们这行的都清楚模型“能跑”和“能打”完全是两码事。同样的训练脚本A 同学一卡一天跑完B 同学三卡三天还在等差距往往不在算法而在工程细节。PyTorch 性能调优这件事很多同学是等项目卡住了才想起来做结果一头扎进 Profiling 日志里翻来覆去越看越迷茫。这篇文章我会把 PyTorch 性能调优的完整思路走一遍围绕 Profiling、torch.compile 和分布式扩展三条主线来讲配合实际踩过的坑和一些可以直接抄走的参数组合希望能帮还在“调优门外”的同学少走几段弯路。这篇学习笔记比较适合两类读者一类是模型已经跑通但训练速度不满意想系统学习如何定位瓶颈的算法工程师另一类是准备把单卡脚本改成多卡分布式训练的工程同学。我会尽量把“为什么这么做”也讲透而不只是丢给你一堆配置项。1. 性能调优的整体思路先测量再动手1.1 别凭感觉猜瓶颈先用数据说话我做性能调优的头一条原则就是不 Profiling 不优化。这句话听起来像废话但实际操作中几乎所有人都会犯“凭经验猜测”的毛病。比如训练变慢了第一反应往往是“是不是并行度不够加卡吧”结果一测GPU 利用率只有 30%瓶颈根本不在算力而在数据加载或者 CPU 端的预处理逻辑。再比如有人一遇到显存不够就想到梯度累积结果 Profiling 一看爆炸显存的其实是某个中间激活值改一下 checkpointing 策略就解决了压根不用动训练策略。所以在优化之前先花十几分钟跑一次完整的 Profiling把 CPU 时间、GPU 时间、显存占用、算子耗时分布、数据加载时间这些数据摆到桌面上。有了这些数据你才知道自己的时间该花在哪个环节而不是像无头苍蝇一样到处瞎调。为什么很多人不愿意先做测量因为 Profiling 看起来有点门槛输出结果密密麻麻不知道看哪里。但现在的 PyTorch Profiler 做得已经相当友好了能自动出表格、火焰图甚至直接输出到 TensorBoard。花十几分钟学会看这些数据后面省下来的时间是以“天”为单位的。1.2 性能问题的三个层面我习惯把 PyTorch 性能问题分成三个层面来排查顺序不能乱第一层是存储和数据流动。包括 DataLoader 的加载速度、预处理是否占用了太多 CPU 时间、数据在 CPU 和 GPU 之间的拷贝是否频繁。很多训练任务慢根本不是算力不够而是“饭”端不上来GPU 在空等数据。第二层是计算效率。包括算子本身的实现效率、是否有大量小算子导致的 kernel launch 开销、是否存在不必要的同步点。这一层的优化通常靠算子融合、把动态图变成静态图、减少 GPU 和 CPU 之间的同步。第三层是扩展效率。单卡已经调得很好了但上多卡之后发现加速比远低于理论值往往是通信开销、负载不均衡、梯度同步策略有问题。这一层需要在分布式层面做手脚。每一层的问题要用不同的工具和手段去解决这也是为什么我不建议一上来就套用网上现成的“性能优化十连”配置。你得先弄清楚自己的问题在哪一层再对症下药。2. Profiling定位瓶颈的第一把手术刀2.1 PyTorch Profiler 的基本用法与关键参数PyTorch 自带的torch.profiler已经是我日常工作里最常用的性能分析工具了。它在 PyTorch 1.8 之后成为官方推荐方案替代了老旧的torch.autograd.profiler用起来非常简单import torch from torch.profiler import profile, ProfilerActivity, record_function def run_training_step(model, batch, optimizer, criterion): optimizer.zero_grad() outputs model(batch) loss criterion(outputs, batch[labels]) loss.backward() optimizer.step() return loss.item() # 用 profiler 包住训练循环 with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3, repeat1), on_trace_readytorch.profiler.tensorboard_trace_handler(./profile_logs), record_shapesTrue, profile_memoryTrue, ) as prof: for step in range(10): loss run_training_step(model, batch, optimizer, criterion) prof.step() # 通知 profiler 一个 step 结束 # 打印统计表 print(prof.key_averages().table(sort_bycuda_time_total, row_limit20))这里有几个参数值得解释一下schedule参数里的wait、warmup、active和repeat是用来控制采样节奏的。wait1表示前 1 个 step 不记录这是为了让模型处于稳定状态warmup1表示接下来 1 个 step 用于预热机制是为了让缓存分配、cuDNN autotune 这些一次性开销先被触发掉active3才是真正采集数据的 3 个 step。这套组合拳打下来拿到的是稳定状态下的性能数据而不是包含各种初始化开销的数据。record_shapesTrue可以记录每个算子的输入张量形状。这个参数在分析内存分配和算子融合效果时非常关键但也有一点代价会让 Profiling 变慢一点。我通常是先不开它跑一遍看整体再单独开它看某个可疑算子的输入形状。profile_memoryTrue会记录内存分配情况定位显存瓶颈时很好用但也会增加开销。实际使用中我一般只在排查显存问题时才打开。on_trace_ready可以指定 trace 处理函数。tensorboard_trace_handler会把结果输出到指定目录之后通过 TensorBoard 的 PyTorch Profiler 插件查看火焰图、算子时间线。这个可视化能力非常强大能看到每个 kernel 在 GPU 上的执行顺序、间隙、以及 CPU 与 GPU 之间的依赖关系。注意Profiling 本身是有开销的尤其是开启record_shapes和profile_memory之后采样期间的速度会比正常训练慢不少。所以不要用 Profiling 阶段的数据去衡量真实训练速度它只用于定位瓶颈。真实速度的评估还是用简单的time.time()包住整个训练循环更靠谱。2.2 三张表看懂输出先看全局再追细节prof.key_averages().table()输出是一张很大的表刚接触的同学往往盯着前几行就懵了。我一般只看三组关键指标第一组是Self CPU time total算子自身的 CPU 耗时不含子算子和Self CUDA time total算子自身的 GPU 耗时。这两个指标能帮你识别最耗时的算子是关注的重点。第二组是CUDA time total和CPU time total的比值。如果所有算子的 GPU 耗时加在一起远小于 CPU 耗时说明 GPU 在等 CPU 喂数据或者 Launch Kernel这种情况大概率有 CPU 端的瓶颈。第三组是Call Stack调用栈和Input Shapes。当定位到某个耗时的算子时点进调用栈看它是从哪个模块发起的这样才能把问题映射到代码层面。看表的顺序也很重要我不建议从耗时最高的算子开始一个一个优化而是先看整体占比结构。比如某个算子占了 60% 的 GPU 时间这通常不是什么问题除非它的理论耗时不该这么高。真正值得警惕的是那种单个耗时占比不大、但数量极多的小算子——它们加起来的总开销非常可观而且往往意味着可以做算子融合来消除。下面这张表我整理了几个关键指标的解读思路方便你对着自己的 Profiling 结果做对照判断关键指标数值偏高意味着什么常见解决方向GPU 利用率~100% 为佳GPU 常处于空闲状态等待数据或同步优化 DataLoader、减少 CPU-GPU 同步CPU 总耗时远高于 CUDA 总耗时CPU 端预处理/数据加载成了瓶颈增加 DataLoader 的num_workers使用pin_memoryKernel Launch 占比高小算子过多launch 开销被放大用torch.compile做算子融合减少小算子数量单个算子耗时异常高可能使用了非最优实现考虑替换为更高效的算子或用 chunk 方式拆解显存峰值接近上限中间激活值或优化器状态占用过大开启 gradient checkpointing、reduce-overhead模式2.3 实际案例一个 Transformer 训练脚本的 Profiling 现象举一个我实际处理过的例子。有个朋友跑来问我他的 BERT 微调脚本在 A100 上 GPU 利用率只有 40% 左右训练 10 万步预计要跑两周急得不行。我让他先跑一次 Profiling结果非常典型CPU 侧耗时几乎等于 GPU 侧耗时而且DataLoader的迭代时间占了 CPU 总耗时的 37%。进一步看时间线GPU 在一段 kernel 执行完之后会有一段很长的空白然后才等来下一批数据。问题立刻清楚了数据加载和预处理跟不上 GPU 的消费速度。再往下细查他的数据管道里有个自定义的 tokenizer 做了非常重的文本清理操作但 DataLoader 的num_workers还停留在默认值 0。也就是说所有预处理都在主进程里跑GPU 只能在旁边干等着。解决方案也很直接data_loader DataLoader( dataset, batch_size32, shuffleTrue, num_workers8, # 从 0 提到 8让子进程负责数据加载 prefetch_factor4, # 每个 worker 预取 4 个 batch pin_memoryTrue, # 锁页内存加速 CPU 到 GPU 的拷贝 )改完再测GPU 利用率直接从 40% 干到了 85% 以上训练时间预计压缩到一周以内。这个案例说明了一个朴素的道理在你往模型结构上折腾之前先老老实实跑一次 Profiling看看瓶颈到底在哪一层。很多时候问题压根出在数据管道而不是模型本身。3. torch.compile一行代码的“静态化”红利3.1 torch.compile 到底做了什么PyTorch 2.0 之后最让人兴奋的改动就是torch.compile。它的原理说起来也不算复杂通过 TorchDynamo 在 Python 字节码层面捕获模型的前向计算图然后用 TorchInductor 把这套计算图编译成高性能的 GPU kernel默认生成 Triton 内核并在这个过程中做算子融合、显存规划、循环优化等一系列自动优化。你可以把它理解成把“逐行解释执行”的 Python 代码变成“整体编译执行”的 C 代码。对于 Python 这种解释型语言来说这个收益是巨大的尤其是当模型里有很多小算子、很多 Python 层循环的时候。用起来确实是“一行代码”的事情import torch model MyModel().cuda() model torch.compile(model) # 这行就够了但这里有个很重要的预期管理torch.compile不是所有模型都能白捡加速的。它最擅长的是那种算子小而多、Python 逻辑重的模型比如基于 Transformer 的各类任务但对于已经高度优化、单算子就是大头比如大矩阵乘法、卷积的模型收益有限甚至可能因为编译开销导致整体变慢。3.2 三种编译模式怎么选torch.compile提供了几种模式我实际用下来比较推荐的是default最均衡的选择。编译时间适中运行速度大部分场景下比 eager 模式快 20%-50%。适合日常开发和验证。reduce-overhead会在 default 的基础上进一步减少 Python 端到 CUDA 端的启动开销对小型 batch 场景尤其有效。代价是会增加显存占用因为它会尽量把中间结果留在 GPU 上。max-autotune编译时会花长时间跑各种 Triton kernel 的 autotune找到当前硬件上最快的实现。训练速度提升最明显但编译时间可能长达十几分钟而且显存占用进一步增加。一般只在最终跑实验、确定不改代码的时候用。我自己的习惯是开发期间用default要出最终训练效果了开一次max-autotune。如果你显存比较紧张可以优先试试reduce-overhead它往往能在显存和速度之间找到一个不错的平衡点。还有dynamicTrue这个参数值得单独说一嘴。默认情况下torch.compile会按照第一次运行时的输入形状完成优化如果后续输入形状变化很大它会重新编译反而变慢。如果你的输入序列长度不固定比如 NLP 任务里的变长 batch建议显式设置dynamicTrue让它提前对动态形状做优化减少反复编译带来的额外开销。3.3 torch.compile 的兼容性“暗坑”torch.compile虽然好用但远没到“一行代码走天下”的程度。以下几个坑是我实际踩过的第一个是动态控制流的兼容问题。Python里常见的if、for、while依赖运行时张量条件、在循环体内动态改变张量形状这些情况 TorchDynamo 可能捕获失败。不过好消息是它有一套 fallback 机制会把无法捕获的部分回退到 eager 模式并打印警告。出现的警告别直接无视最好逐条看有些回退会导致模型核心逻辑没有被编译性能提升会大打折扣。第二个是不支持某些第三方算子。比如一些自定义 CUDA extension、部分 HuggingFace 算子库里的复杂算子都可能触发回退或直接报错。遇到这种情况可以先看报错信息里是否提示“Graph break”然后针对性地把不兼容的子模块从编译范围里排除model torch.compile(model) model.layers.unsupported_module model.layers.unsupported_module.to(memory_formattorch.contiguous_format) # 或者直接对目标子模块“原样保留”只编译其他部分第三个是编译时间和启动延迟。大模型首次编译可能要几分钟这本身不是问题问题在于每次更换机器、更换 GPU 型号之后缓存失效又得重新编译。目前 PyTorch 会在磁盘上缓存编译结果但缓存的有效性依赖于各种环境因素如果频繁切换环境建议把关键运行都固定在同一套镜像或环境里。第四个是显存占用会上升。reduce-overhead和max-autotune模式会倾向于把更多中间结果保存在显存中这在显存紧张的模型中可能直接 OOM。如果碰上了要么切回default模式要么显式开启gradient_checkpointing来换取一部分显存空间。3.4 torch.compile 叠加 Profiling11 2 的效果有人觉得torch.compile之后就不需要 Profiling 了这是天大的误解。恰恰相反编译模式下的 Profiling 价值更高因为它能直接告诉你哪些算子被融合了、哪些算子还残留着。我一般会这样操作先跑一次 eager 模式的 Profiling记录各个算子的耗时占比然后开torch.compile再跑一次同样的 Profiling对比前后算子的耗时变化。如果某类算子耗时在编译后依然占比很高说明它没有被很好地融合你可能需要手动改写这部分逻辑或者调整编译模式。我之前做一个小型语音模型训练时eager模式下LayerNorm占了 GPU 耗时 8%torch.compile之后直接降到 1% 以下因为被融合进了前一个算子。但某个自定义的 attention mask 算子怎么都融合不动一看 Input Shapes发现它在运行时动态生成了 mask 张量导致编译器的图优化无法覆盖。把 mask 生成挪到模型外部、变成静态输入之后编译优化终于追上了这个算子整体速度又快了一截。这就是 Profiling 和torch.compile配合着用的意义编译器能帮你省掉一部分人工优化的功夫但省到什么程度、还剩什么残留必须靠数据来判断。4. 分布式扩展单卡极限之后的必经之路4.1 什么时候该上分布式什么时候不该上很多同学的直觉是“训练速度慢了多搞几张卡”。但实际上分布式训练有启动成本、通信成本、调试成本在模型规模和单卡算力不匹配的情况下盲目上多卡反而可能更慢。我个人的判断标准很简单如果单卡训练时 GPU 利用率已经超过 90%且单卡显存还有富余这时候加卡通常能线性加速收益明显。但如果单卡利用率不到 50%优先解决的是单卡效率问题而不是扩展问题。把单卡从 50% 调到 90%比多上两张卡的效果都好。另一个判断标准是模型显存。如果你的模型已经大到单卡放不下哪怕不用追求极致速度也要考虑分布式训练至少可以用数据并行拆分显存占用或者用模型并行/FSDP 来分配模型参数和梯度。4.2 DDP 的原理与使用要点PyTorch 最常用的分布式训练方案是DistributedDataParallelDDP它的核心思想是每个进程持有一份完整的模型副本各自处理不同的数据 batch前向传播完成后所有进程通过 all-reduce 通信同步梯度保证每个副本都更新到相同的参数。DDP 的使用看起来很简单import torch import torch.distributed as dist import torch.multiprocessing as mp def train_worker(rank, world_size): # 初始化进程组 dist.init_process_group( backendnccl, init_methodtcp://127.0.0.1:23456, rankrank, world_sizeworld_size, ) torch.cuda.set_device(rank) model MyModel().to(rank) model torch.nn.parallel.DistributedDataParallel(model, device_ids[rank]) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) dataset MyDataset() sampler torch.utils.data.distributed.DistributedSampler( dataset, num_replicasworld_size, rankrank ) dataloader DataLoader(dataset, batch_size32, samplersampler) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 保证每个 epoch 数据打乱方式不同 for batch in dataloader: optimizer.zero_grad() outputs model(batch) loss criterion(outputs, batch[labels]) loss.backward() optimizer.step() dist.destroy_process_group() if __name__ __main__: world_size 4 mp.spawn(train_worker, args(world_size,), nprocsworld_size)有几个细节我在实战中反复踩过坑DistributedSampler一定要配合sampler.set_epoch(epoch)使用否则每个 epoch 的数据顺序不变模型容易过拟合到固定的 batch 排列上。batch_size指的是每个 GPU 上的 batch size不是全局 batch size。总 batch size 是单卡 batch_size * world_size。这也意味着如果你原本单卡训练用的 batch size 是 64上 4 卡之后如果每卡还是 64全局 batch size 就变成 256学习率通常要跟着调大一些线性缩放法则否则收敛效果可能变差。DDP 的效率瓶颈往往在梯度同步的通信开销上尤其是卡间通信带宽有限的时候。如果 4 卡加速比远低于 3.5 倍先看NCCL_DEBUGINFO的日志确认通信是否成了瓶颈。4.3 FSDP当模型大到单卡放不下时DDP 虽然好用但每个进程都要持有完整的模型副本。当模型参数优化器状态梯度超过单卡显存时DDP 就无能为力了。这时候就要考虑FullyShardedDataParallelFSDP。FSDP 的核心思路是把模型参数、梯度和优化器状态分片到多个 GPU 上每个 GPU 只保存自己那部分。计算某个算子时需要用到完整的参数再通过通信把分片聚合起来。这个思想和大模型预训练里常见的 ZeRO 策略一脉相承。FSDP 的使用方式也很简单from torch.distributed.fsdp import FullyShardedDataParallel as FSDP model MyBigModel().to(rank) model FSDP( model, sharding_strategytorch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, )sharding_strategy有三个主要选项FULL_SHARD参数、梯度、优化器状态全分片显存最省但通信最重、SHARD_GRAD_OP只分片梯度和优化器状态参数完整复制通信稍轻、NO_SHARD相当于 DDP不分片。实际选型时我的经验是能上单卡尽量单卡单卡放不下且参数在 1B 以下优先考虑SHARD_GRAD_OP参数超过 1B再考虑FULL_SHARD但这时候通信开销会显著拉低训练效率最好配合梯度累积来摊薄通信成本。FSDP 和 DDP 在性能表现上还有个差异FSDP 的通信是穿插在前向反向计算过程中的不是像 DDP 那样在反向结束之后统一同步。所以 FSDP 对计算和通信的重叠要求更高配置不好容易出现“GPU 在等通信”的情况。这就需要借助 Profiling 来看时间线定位是否存在大段的通信等待间隔。4.4 分布式性能排查的常见手段分布式训练出了问题定位起来比单卡复杂一个量级。我常用的排查手段有三板斧第一板斧是看NCCL_DEBUGINFO日志。NCCL 是英伟达多卡通信库打开这个环境变量之后它会输出每个通信算子的耗时、带宽、使用的通信算法等信息。重点看有没有timed out、no peer access、connect to x.x.x.x failed之类的异常。第二板斧是跑一次torch.distributed.all_reduce的纯通信压力测试看多卡之间的实际带宽是否正常。如果通信带宽远低于硬件理论值要么是网络配置问题要么是通信库版本不匹配这和模型代码无关。第三板斧是在多个 worker 上开启 Profiling。分布式场景下不要只在 rank 0 上跑 profiler最好让所有 rank 都输出 trace然后在 TensorBoard 里把时间线对齐看。如果发现某个 rank 的算子执行时间明显偏长大概率是负载不均衡如果所有 rank 都在等待同样长度的通信区间那就是批量太小、通信占比太高。分布式训练还有一个容易被忽略的“减速带”checkpoint 保存。如果每个 epoch 都把模型保存到共享文件系统而模型够大、分布式文件系统带宽不够那么保存阶段可能比训练本身还慢。我的建议是优先保存 rank 0 的模型即可并且可以异步保存不要阻塞训练主循环。5. 常见问题与调优实战速查5.1 问题排查速查表我在项目里积累了一份“踩坑速查表”每次调优遇到问题先对着查一遍解决了不少看起来莫名其妙的问题。这里直接分享出来现象可能原因排查手段与解决方案GPU 利用率长期低于 50%数据加载/预处理瓶颈CPU-GPU 同步过多增加num_workers、pin_memory使用torch.compile减少 launch 开销显存峰值超限中间激活值过多优化器状态过大开启gradient_checkpointing用 FSDP 分片优化器状态降低torch.compile优化等级torch.compile后速度反而更慢动态形状导致频繁重编译图捕获失败回退 eager设置dynamicTrue检查 Graph break 警告排除不兼容子模块多卡加速比远低于显卡数通信等待负载不均衡batch 太小增大单卡 batch开启torch.backends.cudnn.benchmark检查 NCCL 带宽和拓扑DDP 训练 loss 不收敛学习率未随全局 batch size 调整数据采样顺序问题按线性缩放法则上调学习率正确设置sampler.set_epoch(epoch)FSDP 训练时通信开销过大分片策略过重通信没有和计算重叠切换为SHARD_GRAD_OP使用torch.compile让计算和通信更紧凑训练中途报NCCL timeout网络不稳定某进程挂掉跨节点通信配置错误增加 timeout 参数检查防火墙和网络拓扑确保所有进程的 batch 数一致5.2 我最常用的一套调优工作流如果你是新项目我的建议是不要一上来就把所有调优手段全堆上。这样出了问题你都分不清是谁导致的。推荐按照下面的顺序逐步叠加每加一步验证一次效果先用默认配置跑通一个小数据集得到 baseline 数据。Profiling 单卡训练重点解决 GPU 利用率低的问题数据加载、同步、小算子。开启torch.compile(default)对比 Profiling 前后的算子耗时变化。如果模型放不下单卡再上 DDP 或 FSDP从 2 卡开始验证加速比确认无异常再加到预期卡数。训练稳定后再考虑max-autotune、梯度累积、混合精度等进一步优化。这套工作流我用了很多次基本没有翻过车。关键是它保证每一步的变量都足够小出了问题可以快速二分定位。5.3 一个“调优工具箱”级的配置参考最后给大家一套我现在做视觉模型训练时常用的基准配置可以作为你调参的起点。注意这些不是银弹不同任务需要微调但作为参考还是很香的# 训练环境相关 torch.backends.cudnn.benchmark True # 输入 shape 固定时开启自动选最优卷积算法 torch.backends.cuda.matmul.fp16_allow_bf16_reduced_precision_reduction True # 混合精度矩阵乘的精度/速度权衡 # DataLoader dataloader DataLoader( dataset, batch_size64, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, # epoch 间不销毁 worker 进程减少重建开销 ) # 模型 model torch.compile(model, modereduce-overhead) # 训练 scaler torch.cuda.amp.GradScaler() # 配合 amp 自动混合精度混合精度AMP这个点值得单独强调一下。开启torch.cuda.amp.autocast配合GradScaler在大多数模型上能直接带来 1.5-2 倍的提速显存占用也会明显下降。代价是精度可能掉一点但对绝大多数任务来说影响很小属于“性价比”极高的优化。如果你用的是较新版本的 PyTorch比如 2.3 之后AMP 已经默认开启了一些优化但显式控制GradScaler的使用还是能帮你避开一些混合精度下的梯度下溢问题。6. 实操中的心得体会最后分享几个我做性能调优这几年的个人体会不一定系统但都很实用。第一性能调优的回报曲线是阶梯式的。不是说你每投入一小时调优都能稳定换到等比例的加速。很多时候把最明显的那个瓶颈打通之后速度会跳上一个台阶然后进入平台期。这时候不要盲目继续抠细节而是要先问一句现在的速度够用了吗实验迭代的时间成本能不能接受如果已经够用停手是对的继续调优的边际收益可能很低。第二要养成“调参前先备份”的习惯。尤其是跑分布式训练的时候一次错误的 batch size 调整或者学习率调整可能会浪费好几天的实验时间。我一般会给每次实验记录一个配置文件包含模型结构、batch size、学习率、优化器参数、编译模式、混合精度设置。这样即使效果崩了也能快速回到之前的稳定点。第三善用社区经验但别盲信。网上很多性能调优的教程和配置看起来很有道理但硬件环境、模型结构、数据特性不一样效果可能完全相反。我在实操中就遇到过同样一个torch.compile配置在自己的 3090 上加速了 40%换到 A100 上反而慢了 8%。所以任何配置都值得在自己环境里用 Profiling 结果验证一遍。第四别忽略版本的隐性问题。PyTorch 的迭代速度极快同样的代码在不同小版本甚至不同 CUDA 版本下的性能差异可能非常大。如果你遇到奇怪的性能问题先检查版本配套关系说不定只是环境不匹配导致的。更新或回退版本之前务必备份当前可用环境别切完之后发现新的环境还不如旧的。性能调优这件事说难也难说简单也简单。难在它是“系统性工程”涉及数据管道、计算框架、硬件资源、算法结构多个层面说简单是因为只要按照“测量-分析-优化-再测量”的循环一步步来绝大多数性能问题都能被定位和解决。希望这篇笔记能帮你在 PyTorch 性能调优的路上少踩几个坑把宝贵的实验时间花在更有意义的地方。
返回列表