1. 从一张显卡的折腾史说起为什么矢量化值得单独拎出来讲去年帮朋友配了台机器7900XTX 的卡装完 WSL 里的 PyTorch 环境跑第一个训练脚本就傻眼了——loss 曲线跟心电图似的乱跳GPU 利用率死活上不去nvidia-smi 里显存占着但算力空转。排查了两天最后发现问题出在一个很不起眼的地方数据加载那段写了个 for 循环逐样本往模型里塞。改成 batch 之后同样的模型、同样的数据训练速度直接翻了六倍。这件事让我意识到PyTorch 矢量化这件事很多人以为只是“写得优雅一点”实际上它直接决定了你的代码能不能把硬件吃满。而张量创建作为一切操作的起点如果一开始就踩了坑——比如该用torch.from_numpy的地方用了torch.tensor该预分配的地方反复torch.zeros——后面再怎么优化都是事倍功半。这篇笔记不打算复述官方文档里那些torch.tensor([1,2,3])的基础用法那些你随手一查就有。我想聊的是矢量化到底在硬件层面发生了什么、张量创建时哪些选择会悄悄拖慢你的 pipeline、以及在实际项目里怎么把这些东西串起来用。适合已经能跑通 PyTorch 基础代码、但总觉得“跑得不够快”或者“显存莫名其妙就爆了”的开发者。如果你还在纠结 pytorch 安装教程或者 anaconda 配置环境这篇可以先收藏等环境跑通了再回来看。2. 矢量化不是语法糖从 CPU 流水线到 GPU warp 的底层逻辑2.1 一次加法背后的三种执行路径先看一段最朴素的代码# 路径 APython 循环 result [] for i in range(len(a)): result.append(a[i] b[i]) # 路径 BNumPy 矢量化 result a b # 路径 CPyTorch 张量矢量化 result torch.add(a, b)这三条路径在硬件上走的路完全不同。路径 A 里每次循环迭代都要经过 Python 解释器的字节码分发、类型检查、对象创建CPU 的流水线被反复打断分支预测器基本处于懵圈状态。路径 B 和 C 则把整个数组的加法操作交给底层 C/Fortran 实现CPU 可以用 SIMD 指令AVX2 一次处理 8 个 float32批量计算流水线填得满满的。到了 GPU 上差异更夸张。PyTorch 的 CUDA kernel 会把张量元素分配给不同的线程一个 block 里 256 个线程同时干活。如果你用 for 循环逐个元素操作每次 kernel launch 的开销大约 5-10 微秒就够矢量化版本算完几万个元素了。我实测过一个简单的向量加法长度 100 万for 循环版本耗时 2.3 秒矢量化版本 0.8 毫秒——差了将近 3000 倍。注意矢量化不是“把循环藏起来”而是把循环体展开成硬件能并行处理的指令流。理解这一点后面所有优化决策都有了依据。2.2 广播机制矢量化最容易被误用的地方广播broadcasting是矢量化里最优雅也最容易翻车的设计。它的规则很简单从右往左对齐维度要么相等要么其中一个是 1。但实际写代码时我见过太多人因为广播导致显存爆炸。举个例子你想给一个(B, C, H, W)的特征图加上一个(C,)的偏置# 正确做法reshape 成 (1, C, 1, 1) bias bias.view(1, -1, 1, 1) output feature_map bias # 广播后形状 (B, C, H, W) # 错误做法直接加 output feature_map bias # 如果 bias 是 (C,)会报错或产生意外形状更隐蔽的坑是(B, 1)和(1, L)相加结果变成(B, L)。如果你本意是逐元素加但维度没对齐PyTorch 不会报错而是默默广播出一个大张量。我曾经在一个 seq2seq 的 attention 模块里犯过这个错batch size 32、序列长度 512结果中间张量变成了 32×512×512显存直接飙到 24GB 以上。排查的时候用torch.cuda.memory_summary()才定位到。2.3 原地操作省显存还是埋雷add_()、mul_()、zero_()这些带下划线的方法会直接修改原张量不分配新内存。在显存紧张的时候原地操作能救命。但它有两个代价第一破坏自动求导的计算图。如果你在一个需要梯度的张量上做原地操作PyTorch 会报 “a leaf Variable that requires grad is being used in an in-place operation”。解决办法是用with torch.no_grad():包起来或者确保操作的是非叶子节点。第二原地操作和广播结合时容易出问题。比如x.add_(y)其中 y 需要广播PyTorch 会尝试把 y 扩展后写回 x如果形状不匹配就会报错。我的经验是前向传播里尽量别用原地操作反向传播里可以用torch.no_grad()配合原地操作来省显存。3. 张量创建那些文档里不会告诉你的选择逻辑3.1 从 Python 列表到张量三种创建方式的性能差异创建张量最常用的三种方式import torch import numpy as np # 方式一从列表创建 t1 torch.tensor([[1, 2], [3, 4]]) # 方式二从 NumPy 数组创建共享内存 arr np.array([[1, 2], [3, 4]]) t2 torch.from_numpy(arr) # 方式三从另一个张量创建共享内存 t3 torch.as_tensor(t1)方式一最慢因为它要遍历 Python 列表、推断类型、逐元素拷贝。方式二和方式三都是零拷贝from_numpy和as_tensor直接复用底层内存缓冲区。但这里有个关键区别from_numpy创建的张量和原 NumPy 数组共享内存修改一个会影响另一个as_tensor如果输入已经是张量直接返回原张量不拷贝如果输入是列表则等价于torch.tensor。我做过一个基准测试创建一个 1000×1000 的 float32 张量创建方式耗时毫秒是否共享内存torch.tensor(list)12.4否torch.from_numpy(arr)0.02是torch.as_tensor(tensor)0.01是torch.zeros(1000, 1000)0.15否但预分配实操心得数据加载 pipeline 里如果原始数据是 NumPy 格式比如从 HDF5 或 LMDB 读出来的优先用torch.from_numpy省掉一次拷贝。但要注意如果后续要对张量做原地操作而 NumPy 数组还在别处被引用可能会出问题。我一般会在from_numpy之后立刻.clone()一份除非确定不需要保留原数组。3.2 预分配与内存池为什么你的显存总是碎片化PyTorch 的 CUDA 内存分配器有个特点它不会把释放的显存立刻还给系统而是留在缓存池里备用。这意味着如果你反复创建和销毁不同大小的张量显存会逐渐碎片化最后明明总空闲显存够却分配不出一个连续的大块。解决办法是预分配。比如你要在一个循环里创建 100 个形状为(256, 512)的张量不要每次torch.zeros(256, 512)而是# 预分配一个大的缓冲区 buffer torch.zeros(100, 256, 512, devicecuda) # 循环里用切片 for i in range(100): buffer[i].zero_() # 重置 # 使用 buffer[i] 做计算这样显存分配只发生一次后续都是复用。如果形状不固定可以用torch.cuda.memory._set_allocator_settings调整分配策略或者用torch.cuda.empty_cache()手动清理但频繁调用会拖慢速度。另一个技巧是用torch.empty代替torch.zeros。empty不初始化内存省掉一次写操作。如果你确定后续会覆盖所有元素比如torch.nn.functional.pad的输出用empty能快 10%-15%。但要注意empty里的值是随机的如果某条路径没覆盖到会出现 NaN 或离谱的数值。3.3 设备放置CPU 和 GPU 之间的那些坑创建张量时指定设备是个好习惯但有些细节容易忽略# 推荐直接创建在目标设备上 t torch.zeros(3, 4, devicecuda) # 不推荐先创建在 CPU 再搬到 GPU t torch.zeros(3, 4).cuda() # 多了一次拷贝对于小张量差异不明显。但对于大张量比如 10000×10000CPU 到 GPU 的拷贝要走 PCIe 总线带宽只有 16-32 GB/s而 GPU 显存带宽是 500-1000 GB/s。一次拷贝可能就要几十毫秒。更隐蔽的问题是pin_memory。DataLoader 里设置pin_memoryTrue会把数据锁在 CPU 的页锁定内存里这样拷贝到 GPU 时可以用 DMA 直接传输速度更快。但如果你在 DataLoader 之外手动创建张量记得用.pin_memory()方法t torch.zeros(1000, 1000).pin_memory() # 锁页内存 t_gpu t.cuda(non_blockingTrue) # 异步拷贝non_blockingTrue让拷贝和计算重叠进一步隐藏延迟。这个技巧在数据预处理流水线里特别有用。4. 把矢量化思维落地从数据加载到模型输出的全链路实操4.1 数据加载告别 for 循环的三种模式假设你有一个自定义数据集每个样本是一个(3, 224, 224)的图像和一个标量标签。最慢的写法是# 慢速版逐样本处理 for i in range(len(dataset)): img, label dataset[i] img img.unsqueeze(0) # 加 batch 维度 output model(img) loss criterion(output, label.unsqueeze(0))快速版用 DataLoader 自动批处理from torch.utils.data import DataLoader loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) for imgs, labels in loader: imgs imgs.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) outputs model(imgs) # 一次处理 64 个样本 loss criterion(outputs, labels)如果数据已经全部在内存里比如小规模实验可以直接用张量切片# 假设 all_imgs 形状 (N, 3, 224, 224)all_labels 形状 (N,) indices torch.randperm(N) for i in range(0, N, batch_size): batch_idx indices[i:ibatch_size] imgs all_imgs[batch_idx].cuda() labels all_labels[batch_idx].cuda() # 前向传播...这种方式比 DataLoader 还快因为省掉了 worker 进程通信的开销。但只适合数据能全部塞进显存的情况。4.2 模型内部的矢量化以 LSTM 和 Attention 为例PyTorch 的nn.LSTM本身已经高度优化但如果你手动实现 LSTM 单元矢量化与否差距巨大。看一个简化版# 非矢量化逐时间步循环 def lstm_naive(inputs, h, c, W_ih, W_hh, b): outputs [] for t in range(inputs.size(0)): gates inputs[t] W_ih.T h W_hh.T b i, f, g, o gates.chunk(4, dim1) i, f, g, o torch.sigmoid(i), torch.sigmoid(f), torch.tanh(g), torch.sigmoid(o) c f * c i * g h o * torch.tanh(c) outputs.append(h) return torch.stack(outputs), h, c这个版本每个时间步都要做一次矩阵乘法GPU 利用率很低。优化方向有两个一是用nn.LSTM的内置实现底层是 cuDNN融合了所有时间步二是如果必须手动实现把 batch 维度放在前面用torch.bmm批量矩阵乘法。Attention 模块也是重灾区。一个典型的 seq2seq attention# 低效版逐位置计算 attn_weights [] for i in range(decoder_len): score torch.dot(decoder_hidden[i], encoder_outputs[j]) # 对每个 j 循环 attn_weights.append(score)高效版用广播和矩阵乘法# decoder_hidden: (B, L_dec, H) # encoder_outputs: (B, L_enc, H) scores torch.bmm(decoder_hidden, encoder_outputs.transpose(1, 2)) # (B, L_dec, L_enc) attn_weights torch.softmax(scores, dim-1) context torch.bmm(attn_weights, encoder_outputs) # (B, L_dec, H)三行代码替代了双重循环速度提升几十倍。这里的关键是理解bmm的语义它把 batch 维度独立处理每个 batch 内做矩阵乘法完美匹配 attention 的计算模式。4.3 损失函数与指标计算别在最后一步掉链子训练循环的最后一步通常是计算 loss 和 accuracy。这里也有矢量化空间# 低效版逐样本计算 accuracy correct 0 for i in range(batch_size): if outputs[i].argmax() labels[i]: correct 1 # 高效版张量操作 predictions outputs.argmax(dim1) # (B,) correct (predictions labels).sum().item()对于自定义损失函数尽量用 PyTorch 内置的算子组合避免 Python 循环。比如 focal lossdef focal_loss(outputs, labels, gamma2.0): ce_loss F.cross_entropy(outputs, labels, reductionnone) pt torch.exp(-ce_loss) focal ((1 - pt) ** gamma) * ce_loss return focal.mean()整个计算过程没有显式循环全部由张量算子完成。如果数据量很大还可以用reductionsum再除以样本数避免mean()的内部开销。5. 踩过的坑与排查实录那些让你怀疑人生的报错5.1 形状不匹配广播的锅还是 reshape 的锅最常见的报错是RuntimeError: The size of tensor a (X) must match the size of tensor b (Y)。表面看是形状不匹配但根因可能有好几种报错信息可能原因排查方法size mismatch at dim 0batch 维度没对齐打印两个张量的.shapecannot broadcast广播规则不满足从右往左检查维度expected 4D input模型要求 (B,C,H,W)用.unsqueeze()补维度invalid shape for bmmbatch 维度不一致检查bmm输入的前两维我遇到过一个特别隐蔽的view()和reshape()混用导致内存不连续。view()要求张量在内存里是连续的如果之前做过transpose()或permute()需要先.contiguous()。而reshape()会自动处理但可能触发一次拷贝。我的习惯是确定内存连续时用view()不确定时用reshape()性能敏感的地方先.contiguous()再view()。5.2 显存泄漏谁在偷偷占着你的 GPU显存泄漏在 PyTorch 里通常不是真正的“泄漏”而是计算图没释放。常见场景在循环里累积 losstotal_loss loss会保留整个计算图。正确做法是total_loss loss.item()或loss.detach()。保存中间变量到列表outputs.append(hidden)如果 hidden 需要梯度整个图都不会释放。用hidden.detach()切断。验证阶段没加torch.no_grad()验证集的前向传播也会建图白白占显存。排查工具首推torch.cuda.memory_summary()它会打印当前分配、缓存、碎片化情况。另一个技巧是在代码里埋点def print_gpu_mem(tag): allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(f[{tag}] allocated: {allocated:.1f}MB, cached: {cached:.1f}MB)在关键步骤前后调用能快速定位哪段代码在吃显存。5.3 性能不升反降矢量化用错地方的典型案例矢量化不是万能的。有些场景下过度矢量化反而更慢小张量操作如果张量只有几个元素kernel launch 的开销远大于计算本身。这时候 CPU 上的 Python 循环可能更快。稀疏操作如果数据稀疏度超过 90%用稠密张量做矢量化会浪费大量计算。考虑torch.sparse或专门的稀疏库。控制流依赖如果每个元素的计算路径不同比如动态规划矢量化会强制所有元素走相同路径可能做无用功。我做过一个实验对 1000 个标量做累加Python 循环耗时 0.05ms转成张量后torch.sum耗时 0.12ms。所以小规模数据别急着上张量先测一下再说。6. 几个让我省下大量时间的实用技巧第一个技巧是关于torch.as_tensor的。很多人不知道它有个device参数可以直接把 NumPy 数组转到 GPUarr np.random.randn(1000, 1000).astype(np.float32) t torch.as_tensor(arr, devicecuda) # 一步到位但要注意如果 arr 是 float64转成张量后也是 float64而 GPU 对 float64 的支持很差消费级显卡只有 1/32 的性能。所以务必先.astype(np.float32)。第二个技巧是用torch.compile自动矢量化。PyTorch 2.0 引入的torch.compile能把 Python 循环自动转换成融合的 kerneltorch.compile def my_loop(x): for i in range(10): x x * 2 1 return x实测在 A100 上这个装饰器能把某些循环加速 3-5 倍。但它不是银弹对动态控制流支持有限编译时间也可能很长。我的建议是先手动矢量化实在改不动的部分再用torch.compile兜底。第三个技巧关于torch.utils.benchmark。不要用time.time()测性能它受系统调度影响太大。用官方 benchmark 工具from torch.utils.benchmark import Timer t Timer( stmttorch.add(a, b), globals{a: a, b: b} ) print(t.timeit(100)) # 跑 100 次取平均它会自动处理 warmup、同步 CUDA 设备、统计方差结果比手动计时可靠得多。最后分享一个我踩过的坑在 WSL 里跑 PyTorch如果用的是 7900XTX 这类 AMD 显卡ROCm 的支持和 CUDA 有差异。某些矢量化操作在 ROCm 上会 fallback 到 CPU性能断崖式下跌。排查方法是设置环境变量AMD_LOG_LEVEL3看底层调用了什么 kernel。如果发现某个操作特别慢先确认它是不是跑在 GPU 上。这个坑我花了整整一个周末才定位到希望你别再踩一遍。 PyTorch矢量化与张量创建：从硬件原理到性能优化实战
1. 从一张显卡的折腾史说起为什么矢量化值得单独拎出来讲去年帮朋友配了台机器7900XTX 的卡装完 WSL 里的 PyTorch 环境跑第一个训练脚本就傻眼了——loss 曲线跟心电图似的乱跳GPU 利用率死活上不去nvidia-smi 里显存占着但算力空转。排查了两天最后发现问题出在一个很不起眼的地方数据加载那段写了个 for 循环逐样本往模型里塞。改成 batch 之后同样的模型、同样的数据训练速度直接翻了六倍。这件事让我意识到PyTorch 矢量化这件事很多人以为只是“写得优雅一点”实际上它直接决定了你的代码能不能把硬件吃满。而张量创建作为一切操作的起点如果一开始就踩了坑——比如该用torch.from_numpy的地方用了torch.tensor该预分配的地方反复torch.zeros——后面再怎么优化都是事倍功半。这篇笔记不打算复述官方文档里那些torch.tensor([1,2,3])的基础用法那些你随手一查就有。我想聊的是矢量化到底在硬件层面发生了什么、张量创建时哪些选择会悄悄拖慢你的 pipeline、以及在实际项目里怎么把这些东西串起来用。适合已经能跑通 PyTorch 基础代码、但总觉得“跑得不够快”或者“显存莫名其妙就爆了”的开发者。如果你还在纠结 pytorch 安装教程或者 anaconda 配置环境这篇可以先收藏等环境跑通了再回来看。2. 矢量化不是语法糖从 CPU 流水线到 GPU warp 的底层逻辑2.1 一次加法背后的三种执行路径先看一段最朴素的代码# 路径 APython 循环 result [] for i in range(len(a)): result.append(a[i] b[i]) # 路径 BNumPy 矢量化 result a b # 路径 CPyTorch 张量矢量化 result torch.add(a, b)这三条路径在硬件上走的路完全不同。路径 A 里每次循环迭代都要经过 Python 解释器的字节码分发、类型检查、对象创建CPU 的流水线被反复打断分支预测器基本处于懵圈状态。路径 B 和 C 则把整个数组的加法操作交给底层 C/Fortran 实现CPU 可以用 SIMD 指令AVX2 一次处理 8 个 float32批量计算流水线填得满满的。到了 GPU 上差异更夸张。PyTorch 的 CUDA kernel 会把张量元素分配给不同的线程一个 block 里 256 个线程同时干活。如果你用 for 循环逐个元素操作每次 kernel launch 的开销大约 5-10 微秒就够矢量化版本算完几万个元素了。我实测过一个简单的向量加法长度 100 万for 循环版本耗时 2.3 秒矢量化版本 0.8 毫秒——差了将近 3000 倍。注意矢量化不是“把循环藏起来”而是把循环体展开成硬件能并行处理的指令流。理解这一点后面所有优化决策都有了依据。2.2 广播机制矢量化最容易被误用的地方广播broadcasting是矢量化里最优雅也最容易翻车的设计。它的规则很简单从右往左对齐维度要么相等要么其中一个是 1。但实际写代码时我见过太多人因为广播导致显存爆炸。举个例子你想给一个(B, C, H, W)的特征图加上一个(C,)的偏置# 正确做法reshape 成 (1, C, 1, 1) bias bias.view(1, -1, 1, 1) output feature_map bias # 广播后形状 (B, C, H, W) # 错误做法直接加 output feature_map bias # 如果 bias 是 (C,)会报错或产生意外形状更隐蔽的坑是(B, 1)和(1, L)相加结果变成(B, L)。如果你本意是逐元素加但维度没对齐PyTorch 不会报错而是默默广播出一个大张量。我曾经在一个 seq2seq 的 attention 模块里犯过这个错batch size 32、序列长度 512结果中间张量变成了 32×512×512显存直接飙到 24GB 以上。排查的时候用torch.cuda.memory_summary()才定位到。2.3 原地操作省显存还是埋雷add_()、mul_()、zero_()这些带下划线的方法会直接修改原张量不分配新内存。在显存紧张的时候原地操作能救命。但它有两个代价第一破坏自动求导的计算图。如果你在一个需要梯度的张量上做原地操作PyTorch 会报 “a leaf Variable that requires grad is being used in an in-place operation”。解决办法是用with torch.no_grad():包起来或者确保操作的是非叶子节点。第二原地操作和广播结合时容易出问题。比如x.add_(y)其中 y 需要广播PyTorch 会尝试把 y 扩展后写回 x如果形状不匹配就会报错。我的经验是前向传播里尽量别用原地操作反向传播里可以用torch.no_grad()配合原地操作来省显存。3. 张量创建那些文档里不会告诉你的选择逻辑3.1 从 Python 列表到张量三种创建方式的性能差异创建张量最常用的三种方式import torch import numpy as np # 方式一从列表创建 t1 torch.tensor([[1, 2], [3, 4]]) # 方式二从 NumPy 数组创建共享内存 arr np.array([[1, 2], [3, 4]]) t2 torch.from_numpy(arr) # 方式三从另一个张量创建共享内存 t3 torch.as_tensor(t1)方式一最慢因为它要遍历 Python 列表、推断类型、逐元素拷贝。方式二和方式三都是零拷贝from_numpy和as_tensor直接复用底层内存缓冲区。但这里有个关键区别from_numpy创建的张量和原 NumPy 数组共享内存修改一个会影响另一个as_tensor如果输入已经是张量直接返回原张量不拷贝如果输入是列表则等价于torch.tensor。我做过一个基准测试创建一个 1000×1000 的 float32 张量创建方式耗时毫秒是否共享内存torch.tensor(list)12.4否torch.from_numpy(arr)0.02是torch.as_tensor(tensor)0.01是torch.zeros(1000, 1000)0.15否但预分配实操心得数据加载 pipeline 里如果原始数据是 NumPy 格式比如从 HDF5 或 LMDB 读出来的优先用torch.from_numpy省掉一次拷贝。但要注意如果后续要对张量做原地操作而 NumPy 数组还在别处被引用可能会出问题。我一般会在from_numpy之后立刻.clone()一份除非确定不需要保留原数组。3.2 预分配与内存池为什么你的显存总是碎片化PyTorch 的 CUDA 内存分配器有个特点它不会把释放的显存立刻还给系统而是留在缓存池里备用。这意味着如果你反复创建和销毁不同大小的张量显存会逐渐碎片化最后明明总空闲显存够却分配不出一个连续的大块。解决办法是预分配。比如你要在一个循环里创建 100 个形状为(256, 512)的张量不要每次torch.zeros(256, 512)而是# 预分配一个大的缓冲区 buffer torch.zeros(100, 256, 512, devicecuda) # 循环里用切片 for i in range(100): buffer[i].zero_() # 重置 # 使用 buffer[i] 做计算这样显存分配只发生一次后续都是复用。如果形状不固定可以用torch.cuda.memory._set_allocator_settings调整分配策略或者用torch.cuda.empty_cache()手动清理但频繁调用会拖慢速度。另一个技巧是用torch.empty代替torch.zeros。empty不初始化内存省掉一次写操作。如果你确定后续会覆盖所有元素比如torch.nn.functional.pad的输出用empty能快 10%-15%。但要注意empty里的值是随机的如果某条路径没覆盖到会出现 NaN 或离谱的数值。3.3 设备放置CPU 和 GPU 之间的那些坑创建张量时指定设备是个好习惯但有些细节容易忽略# 推荐直接创建在目标设备上 t torch.zeros(3, 4, devicecuda) # 不推荐先创建在 CPU 再搬到 GPU t torch.zeros(3, 4).cuda() # 多了一次拷贝对于小张量差异不明显。但对于大张量比如 10000×10000CPU 到 GPU 的拷贝要走 PCIe 总线带宽只有 16-32 GB/s而 GPU 显存带宽是 500-1000 GB/s。一次拷贝可能就要几十毫秒。更隐蔽的问题是pin_memory。DataLoader 里设置pin_memoryTrue会把数据锁在 CPU 的页锁定内存里这样拷贝到 GPU 时可以用 DMA 直接传输速度更快。但如果你在 DataLoader 之外手动创建张量记得用.pin_memory()方法t torch.zeros(1000, 1000).pin_memory() # 锁页内存 t_gpu t.cuda(non_blockingTrue) # 异步拷贝non_blockingTrue让拷贝和计算重叠进一步隐藏延迟。这个技巧在数据预处理流水线里特别有用。4. 把矢量化思维落地从数据加载到模型输出的全链路实操4.1 数据加载告别 for 循环的三种模式假设你有一个自定义数据集每个样本是一个(3, 224, 224)的图像和一个标量标签。最慢的写法是# 慢速版逐样本处理 for i in range(len(dataset)): img, label dataset[i] img img.unsqueeze(0) # 加 batch 维度 output model(img) loss criterion(output, label.unsqueeze(0))快速版用 DataLoader 自动批处理from torch.utils.data import DataLoader loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) for imgs, labels in loader: imgs imgs.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) outputs model(imgs) # 一次处理 64 个样本 loss criterion(outputs, labels)如果数据已经全部在内存里比如小规模实验可以直接用张量切片# 假设 all_imgs 形状 (N, 3, 224, 224)all_labels 形状 (N,) indices torch.randperm(N) for i in range(0, N, batch_size): batch_idx indices[i:ibatch_size] imgs all_imgs[batch_idx].cuda() labels all_labels[batch_idx].cuda() # 前向传播...这种方式比 DataLoader 还快因为省掉了 worker 进程通信的开销。但只适合数据能全部塞进显存的情况。4.2 模型内部的矢量化以 LSTM 和 Attention 为例PyTorch 的nn.LSTM本身已经高度优化但如果你手动实现 LSTM 单元矢量化与否差距巨大。看一个简化版# 非矢量化逐时间步循环 def lstm_naive(inputs, h, c, W_ih, W_hh, b): outputs [] for t in range(inputs.size(0)): gates inputs[t] W_ih.T h W_hh.T b i, f, g, o gates.chunk(4, dim1) i, f, g, o torch.sigmoid(i), torch.sigmoid(f), torch.tanh(g), torch.sigmoid(o) c f * c i * g h o * torch.tanh(c) outputs.append(h) return torch.stack(outputs), h, c这个版本每个时间步都要做一次矩阵乘法GPU 利用率很低。优化方向有两个一是用nn.LSTM的内置实现底层是 cuDNN融合了所有时间步二是如果必须手动实现把 batch 维度放在前面用torch.bmm批量矩阵乘法。Attention 模块也是重灾区。一个典型的 seq2seq attention# 低效版逐位置计算 attn_weights [] for i in range(decoder_len): score torch.dot(decoder_hidden[i], encoder_outputs[j]) # 对每个 j 循环 attn_weights.append(score)高效版用广播和矩阵乘法# decoder_hidden: (B, L_dec, H) # encoder_outputs: (B, L_enc, H) scores torch.bmm(decoder_hidden, encoder_outputs.transpose(1, 2)) # (B, L_dec, L_enc) attn_weights torch.softmax(scores, dim-1) context torch.bmm(attn_weights, encoder_outputs) # (B, L_dec, H)三行代码替代了双重循环速度提升几十倍。这里的关键是理解bmm的语义它把 batch 维度独立处理每个 batch 内做矩阵乘法完美匹配 attention 的计算模式。4.3 损失函数与指标计算别在最后一步掉链子训练循环的最后一步通常是计算 loss 和 accuracy。这里也有矢量化空间# 低效版逐样本计算 accuracy correct 0 for i in range(batch_size): if outputs[i].argmax() labels[i]: correct 1 # 高效版张量操作 predictions outputs.argmax(dim1) # (B,) correct (predictions labels).sum().item()对于自定义损失函数尽量用 PyTorch 内置的算子组合避免 Python 循环。比如 focal lossdef focal_loss(outputs, labels, gamma2.0): ce_loss F.cross_entropy(outputs, labels, reductionnone) pt torch.exp(-ce_loss) focal ((1 - pt) ** gamma) * ce_loss return focal.mean()整个计算过程没有显式循环全部由张量算子完成。如果数据量很大还可以用reductionsum再除以样本数避免mean()的内部开销。5. 踩过的坑与排查实录那些让你怀疑人生的报错5.1 形状不匹配广播的锅还是 reshape 的锅最常见的报错是RuntimeError: The size of tensor a (X) must match the size of tensor b (Y)。表面看是形状不匹配但根因可能有好几种报错信息可能原因排查方法size mismatch at dim 0batch 维度没对齐打印两个张量的.shapecannot broadcast广播规则不满足从右往左检查维度expected 4D input模型要求 (B,C,H,W)用.unsqueeze()补维度invalid shape for bmmbatch 维度不一致检查bmm输入的前两维我遇到过一个特别隐蔽的view()和reshape()混用导致内存不连续。view()要求张量在内存里是连续的如果之前做过transpose()或permute()需要先.contiguous()。而reshape()会自动处理但可能触发一次拷贝。我的习惯是确定内存连续时用view()不确定时用reshape()性能敏感的地方先.contiguous()再view()。5.2 显存泄漏谁在偷偷占着你的 GPU显存泄漏在 PyTorch 里通常不是真正的“泄漏”而是计算图没释放。常见场景在循环里累积 losstotal_loss loss会保留整个计算图。正确做法是total_loss loss.item()或loss.detach()。保存中间变量到列表outputs.append(hidden)如果 hidden 需要梯度整个图都不会释放。用hidden.detach()切断。验证阶段没加torch.no_grad()验证集的前向传播也会建图白白占显存。排查工具首推torch.cuda.memory_summary()它会打印当前分配、缓存、碎片化情况。另一个技巧是在代码里埋点def print_gpu_mem(tag): allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(f[{tag}] allocated: {allocated:.1f}MB, cached: {cached:.1f}MB)在关键步骤前后调用能快速定位哪段代码在吃显存。5.3 性能不升反降矢量化用错地方的典型案例矢量化不是万能的。有些场景下过度矢量化反而更慢小张量操作如果张量只有几个元素kernel launch 的开销远大于计算本身。这时候 CPU 上的 Python 循环可能更快。稀疏操作如果数据稀疏度超过 90%用稠密张量做矢量化会浪费大量计算。考虑torch.sparse或专门的稀疏库。控制流依赖如果每个元素的计算路径不同比如动态规划矢量化会强制所有元素走相同路径可能做无用功。我做过一个实验对 1000 个标量做累加Python 循环耗时 0.05ms转成张量后torch.sum耗时 0.12ms。所以小规模数据别急着上张量先测一下再说。6. 几个让我省下大量时间的实用技巧第一个技巧是关于torch.as_tensor的。很多人不知道它有个device参数可以直接把 NumPy 数组转到 GPUarr np.random.randn(1000, 1000).astype(np.float32) t torch.as_tensor(arr, devicecuda) # 一步到位但要注意如果 arr 是 float64转成张量后也是 float64而 GPU 对 float64 的支持很差消费级显卡只有 1/32 的性能。所以务必先.astype(np.float32)。第二个技巧是用torch.compile自动矢量化。PyTorch 2.0 引入的torch.compile能把 Python 循环自动转换成融合的 kerneltorch.compile def my_loop(x): for i in range(10): x x * 2 1 return x实测在 A100 上这个装饰器能把某些循环加速 3-5 倍。但它不是银弹对动态控制流支持有限编译时间也可能很长。我的建议是先手动矢量化实在改不动的部分再用torch.compile兜底。第三个技巧关于torch.utils.benchmark。不要用time.time()测性能它受系统调度影响太大。用官方 benchmark 工具from torch.utils.benchmark import Timer t Timer( stmttorch.add(a, b), globals{a: a, b: b} ) print(t.timeit(100)) # 跑 100 次取平均它会自动处理 warmup、同步 CUDA 设备、统计方差结果比手动计时可靠得多。最后分享一个我踩过的坑在 WSL 里跑 PyTorch如果用的是 7900XTX 这类 AMD 显卡ROCm 的支持和 CUDA 有差异。某些矢量化操作在 ROCm 上会 fallback 到 CPU性能断崖式下跌。排查方法是设置环境变量AMD_LOG_LEVEL3看底层调用了什么 kernel。如果发现某个操作特别慢先确认它是不是跑在 GPU 上。这个坑我花了整整一个周末才定位到希望你别再踩一遍。