PyTorch矢量化与张量创建：从循环到批量运算的性能跃迁

1. 从一次踩坑说起为什么矢量化值得单独记笔记刚接触 PyTorch 那会儿我写训练循环的习惯跟写纯 Python 没两样——一个样本一个样本地喂一层一层地手写 for。跑 MNIST 这种小数据集还能忍等到换成几万条文本、几百维特征的业务数据GPU 利用率常年趴在 20% 以下风扇都不怎么转训练一轮要等一炷香。后来被一位做推理优化的老哥点了一句“你这代码根本没矢量化GPU 在等你喂饭。”回去把循环改成整批张量运算同样的模型、同样的数据单轮耗时从 40 多秒掉到 6 秒出头。那次之后我才真正意识到PyTorch 的矢量化不是可选项而是默认姿势而矢量化的一切都建立在张量这个数据结构上。这篇笔记就是围绕“矢量化”和“张量创建”这两件事展开的。它适合刚装完 PyTorch、能跑通官方 MNIST 例子但说不清torch.tensor和torch.Tensor区别的人也适合写了半年模型、想搞清楚为什么自己的代码慢、显存还爆的人。我会把张量创建的各种入口、矢量化背后的原理、实际写代码时容易踩的坑以及排查问题的思路都摊开讲。不堆公式尽量用生活化的类比把“为什么这么写”讲透让你看完能直接改自己手里的代码。需要先说明一点下面涉及的具体 API 行为、参数默认值都是基于 PyTorch 1.x 到 2.x 主流版本的常见实践不同小版本可能有细微差异动手前建议用torch.__version__确认一下自己的环境。另外本文不涉及任何环境安装的特殊渠道安装部分只讲通用思路。2. 张量到底是什么把标量、向量、矩阵串成一条线2.1 从 Python 列表到张量的思维转变很多人第一次看到张量会懵这不就是个多维数组吗跟 NumPy 有啥区别区别大了。你可以把张量理解成“带自动求导和硬件加速能力的多维数组”。普通 Python 列表是散装零件NumPy 数组是装好的积木而张量是积木加上一套“能记住自己怎么被拼出来”的说明书——这套说明书就是计算图和 autograd。举个直观的例子。假设你要算y x^2 3xx 是一个长度为 100 万的向量。用 Python 列表你得写循环一百万个元素挨个算用 NumPy 你能一行x**2 3*x搞定但求导还得自己推公式用张量你同样一行搞定而且y.backward()之后x.grad直接就是2x 3。这就是张量相对前两者的核心增量批量运算 自动微分 设备无关。从维度上看张量是一条连续的谱系0 维是标量一个数1 维是向量2 维是矩阵3 维及以上统称高阶张量。图像数据通常是 4 维(N, C, H, W)即批量大小、通道数、高、宽文本序列经过 embedding 后常见 3 维(N, L, D)。理解这条谱系很重要因为后面所有矢量化操作本质都是在某个维度上做“批量”。2.2 张量和向量、矢量的区别到底在哪热搜里反复出现“张量和向量的区别”“张量和矢量有什么区别”说明这是新手最容易绕晕的点。用一句话概括向量是张量在 1 维时的特例矢量是物理/数学语境下对向量的另一种叫法而张量是更一般的框架。标量只有大小0 维比如温度 25 度。向量矢量有大小有方向1 维比如速度 (3, 4)。矩阵2 维比如一张灰度图。张量任意维度是上面这些的统一容器。所以“张量和矢量的区别”这个问题本身有点错位——矢量就是 1 维张量二者不是并列关系而是包含关系。真正需要区分的是“张量”和“NumPy 数组”前者能上 GPU、能求导、能参与计算图后者不行。搞清楚这一点后面选 API 时就不会纠结。2.3 为什么矢量化能带来数量级的加速矢量化vectorization的核心思想是把逐元素的循环改写成对整个数组的批量操作。为什么快三个层面的原因。第一Python 解释器的循环开销极大。每执行一次for迭代都要做类型检查、引用计数、字节码分派这些开销跟实际计算量比可能高几十倍。矢量化把循环下沉到 C/C 甚至 CUDA 层面Python 只负责发起一次调用。第二底层库如 MKL、cuBLAS对批量运算做了深度优化能利用 SIMD 指令一次处理多个数据还能更好地命中缓存。第三GPU 天生适合并行。一个 4096 维的向量加法CPU 可能分几个批次串行算GPU 可以几千个核心同时开工。但前提是你得把数据一次性交给它而不是一个元素一个元素地传——后者会因为频繁的主机-设备通信把优势全吃掉。我做过一个粗糙的对比对长度为 1000 万的向量做逐元素平方纯 Python 循环大约 3 秒NumPy 约 0.02 秒PyTorch CPU 张量约 0.03 秒PyTorch GPU 张量约 0.005 秒不含传输。差距是三个数量级。这就是为什么“能矢量化就绝不写循环”应该成为肌肉记忆。3. 张量创建的几条路torch.tensor 与 torch.Tensor 的恩怨3.1 torch.tensor推荐的默认入口torch.tensor(data)是我日常用得最多的创建方式。它接收 Python 列表、NumPy 数组、标量等返回一个新的张量并且会根据输入数据自动推断 dtype。比如传整数列表得到int64传浮点列表得到float32。import torch a torch.tensor([1, 2, 3]) # dtypetorch.int64 b torch.tensor([1.0, 2.0, 3.0]) # dtypetorch.float32 c torch.tensor([[1, 2], [3, 4]]) # 2x2 int64它的一个关键特性是总是拷贝数据。你传进去的列表或数组跟返回的张量不再共享内存改一个不影响另一个。这在需要数据隔离时是优点在追求极致内存时是缺点。torch.tensor还支持几个常用参数dtype显式指定类型device指定设备requires_grad决定是否追踪梯度。我一般会显式写dtypetorch.float32因为默认推断有时会给出float64而 GPU 上float64又慢又占显存是个隐形坑。3.2 torch.Tensor历史遗留的构造函数torch.Tensor(data)是更老的写法它其实是torch.FloatTensor的别名。这意味着不管你传什么它默认产出 float32而且对输入的处理规则跟torch.tensor不完全一致。d torch.Tensor([1, 2, 3]) # dtypetorch.float32注意不是 int更坑的是torch.Tensor(3, 4)这种传多个整数的写法会创建一个 3x4 的未初始化张量里面的值是内存里的垃圾数据不是零。新手如果误以为它创建的是全零张量训练时看到 loss 是 NaN 会一脸懵。我的建议很直接新代码一律用torch.tensor除非你在维护老项目或者明确需要 float32 默认行为。两者性能差异可以忽略但可读性和可预测性差很多。3.3 特殊张量的创建zeros、ones、randn、arange除了从数据创建PyTorch 提供了一整套“按形状造张量”的工厂函数这些在初始化权重、构造 mask、生成位置编码时天天用。函数作用典型场景torch.zeros(shape)全零初始化 bias、masktorch.ones(shape)全一构造权重掩码torch.full(shape, v)全为 v填充常量torch.randn(shape)标准正态权重初始化torch.rand(shape)[0,1) 均匀随机采样torch.arange(s, e, step)等差数列位置索引torch.linspace(s, e, n)等间隔 n 个点插值、坐标网格torch.eye(n)单位矩阵线性代数这里有个细节值得说torch.zeros和torch.ones是真正把内存清零/置一的而torch.empty只分配内存不初始化速度最快但内容随机。如果你打算立刻覆盖所有元素用empty能省一点时间否则老老实实用zeros。3.4 从 NumPy 转换from_numpy 与 as_tensor 的取舍实际项目里数据往往先经过 NumPy 或 pandas 处理再喂给 PyTorch。这时有两个常用入口torch.from_numpy(ndarray)共享内存改张量会影响原数组反之亦然。零拷贝快但危险。torch.as_tensor(data)如果输入已经是同 dtype 同设备的张量直接返回否则拷贝。相对安全。import numpy as np arr np.array([1.0, 2.0, 3.0]) t1 torch.from_numpy(arr) # 共享内存 t2 torch.as_tensor(arr) # 可能共享取决于输入我踩过的坑用from_numpy拿到张量后在训练循环里对它做了原地操作比如t 1结果原始 NumPy 数组被悄悄改了导致下一轮 epoch 的数据全错。排查了半天才发现是共享内存惹的祸。所以除非你明确知道自己在做什么优先用torch.tensor或torch.as_tensor做拷贝。4. 矢量化实战把循环改写成张量运算4.1 一个真实的慢代码案例假设有个需求给一批样本计算加权得分。每个样本有 128 维特征权重也是 128 维最后要加上一个偏置。新手写法可能是这样def slow_score(features, weights, bias): results [] for i in range(len(features)): s 0.0 for j in range(len(weights)): s features[i][j] * weights[j] results.append(s bias) return results1000 个样本、128 维特征这段代码在 CPU 上要跑将近 0.5 秒。改成矢量化def fast_score(features, weights, bias): f torch.tensor(features) # (N, D) w torch.tensor(weights) # (D,) return f w bias # (N,)同样的数据耗时降到 0.3 毫秒左右快了三个数量级。是矩阵乘法运算符(N, D) (D,)得到(N,)一步到位。这就是矢量化的威力把两层嵌套循环压缩成一次矩阵乘法。4.2 广播机制矢量化的隐形推手矢量化能写得这么简洁广播broadcasting功不可没。广播的规则是从最后一个维度开始对齐维度相等或其中一个为 1 就能运算为 1 的维度会被“拉伸”到匹配。a torch.randn(3, 1) # (3, 1) b torch.randn(1, 4) # (1, 4) c a b # (3, 4)这里 a 沿列方向复制 4 次b 沿行方向复制 3 次得到 3x4 的结果。广播不实际复制内存只是逻辑上扩展所以非常高效。但广播也是 bug 温床。比如你想算(N, D)和(N,)的逐样本点积直接写x * y会得到(N, N)而不是(N,)因为 y 被当成(1, N)广播了。正确写法是(x * y).sum(dim1)或者torch.einsum(nd,n-n, x, y)。我建议在形状不确定时随手print(x.shape)比事后 debug 省事得多。4.3 用 einsum 表达复杂矢量化逻辑当维度多起来普通矩阵乘法不够用torch.einsum就是救星。它用一套类似爱因斯坦求和约定的字符串描述运算可读性极强。# 批量矩阵乘法: (B, N, M) (B, M, K) - (B, N, K) out torch.einsum(bnm,bmk-bnk, A, B) # 逐样本点积: (N, D) 和 (N, D) - (N,) dot torch.einsum(nd,nd-n, X, Y) # 注意力打分: (B, H, L, D) 和 (B, H, S, D) - (B, H, L, S) scores torch.einsum(bhld,bhsd-bhls, Q, K)einsum 的好处是一次写清楚输入输出维度关系不用记 transpose、reshape 的顺序。代价是它不一定总能选到最优的底层实现极端性能场景下可能需要手写展开。但日常开发einsum 的可维护性远超一堆 permute。4.4 原地操作省内存但别乱用PyTorch 里带下划线后缀的方法如add_、mul_、zero_是原地操作直接改张量本身不新建对象。在显存紧张时能省不少内存。x torch.randn(1000, 1000) x.add_(1) # 原地加x 被修改 y x.add(1) # 非原地返回新张量但原地操作有两个雷区。第一会破坏 autograd 的计算图。如果 x 是需要求导的叶子节点原地修改后反向传播可能报错或给出错误梯度。第二共享内存的张量会被连带修改前面from_numpy的例子就是教训。我的原则是训练循环里涉及梯度的张量一律不用原地操作纯推理或预处理阶段可以放心用。5. 设备、dtype 与性能矢量化之外的三个关键变量5.1 CPU 与 GPU 之间的数据搬运矢量化写对了但如果数据在 CPU 和 GPU 之间反复横跳性能照样上不去。核心原则让数据和模型待在同一个设备上搬运只做一次。device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) data data.to(device) # 一次性搬过去 for batch in loader: batch batch.to(device, non_blockingTrue) out model(batch)non_blockingTrue配合pin_memoryTrue的 DataLoader能让数据搬运和计算重叠进一步压榨吞吐。我实测在图像任务上这个小改动能带来 10% 到 20% 的提速。另外.to(device)如果源和目标设备相同会直接返回原张量不拷贝所以不用担心重复调用有开销。5.2 dtype 选择float32、float16 与 bfloat16默认的float32是通用选择但不同场景有更优解。float64科学计算需要高精度时用GPU 上性能差显存翻倍深度学习基本不用。float16混合精度训练用显存减半、速度提升但数值范围小容易溢出需要配合 loss scaling。bfloat16动态范围跟 float32 一样精度略低在支持它的硬件上做混合精度更省心。with torch.autocast(device_typecuda, dtypetorch.float16): out model(batch) loss criterion(out, target)混合精度不是万能药。小模型、小批量下收益有限反而可能因为类型转换引入额外开销。我一般先跑 baseline确认是显存或算力瓶颈后再上混合精度。5.3 内存布局contiguous 与 view 的关系张量在内存里是一维连续存储的多维形状只是“视图”。view要求张量内存连续reshape则会在必要时自动拷贝。转置、permute 之后张量往往变得不连续这时调view会报错。x torch.randn(4, 6) y x.t() # 转置不连续 # y.view(24) # 报错 y.reshape(24) # OK内部会拷贝 y.contiguous().view(24) # 显式连续化后再 view矢量化代码里频繁 reshape 时这个细节很容易被忽略。我的习惯是能用 reshape 就用 reshape除非我确定内存连续且想避免拷贝。6. 常见问题与排查技巧实录6.1 形状不匹配最常见的报错来源矩阵乘法报size mismatch、广播结果跟预期不符九成是形状问题。排查步骤我总结成三步打印所有参与运算张量的.shape。从右往左对齐维度检查是否满足“相等或为 1”。用unsqueeze或view显式补齐维度别依赖隐式广播。x torch.randn(32, 128) w torch.randn(128, 64) print(x.shape, w.shape) # 确认 out x w # (32, 64)提示调试阶段可以在关键运算前加一行assert比如assert x.shape[-1] w.shape[0]把错误挡在发生之前。6.2 梯度为 None 或报“leaf variable”错误如果loss.backward()后某个参数的.grad是 None常见原因有三个该参数requires_gradFalse该参数没有参与前向计算中间用了原地操作把计算图切断。逐个检查即可。“a leaf Variable that requires grad is being used in an in-place operation”这个报错就是原地操作踩了叶子节点。解决办法是改用非原地版本或者用with torch.no_grad():包住不需要梯度的部分。6.3 显存溢出OOM的排查顺序OOM 是矢量化写过头时的常见后果——一次把整个数据集塞进显存谁都扛不住。排查顺序建议步骤检查项处理方式1batch size 是否过大减半再试2是否有中间张量未释放用deltorch.cuda.empty_cache()3是否在循环里累积了计算图用.item()或.detach()断开4是否用了 float64改 float325是否开了混合精度开启 autocast最容易忽略的是第 3 条在训练循环里写total_loss loss如果 loss 还挂着计算图每轮都会累积显存线性增长。正确写法是total_loss loss.item()。6.4 矢量化后结果对不上数值精度陷阱把循环改成批量运算后有时结果会有微小差异这是浮点累加顺序不同导致的通常可接受。但如果差异很大要检查是不是广播方向搞反了或者 dtype 在中间被悄悄提升/截断。a torch.tensor([1e20, 1.0, -1e20]) print(a.sum()) # 可能得到 0 而不是 1.0浮点精度问题遇到这种情况可以用float64做一次对照计算确认是精度问题还是逻辑问题。6.5 独家避坑清单创建张量时永远显式写dtype别赌默认推断。新代码用torch.tensor别用torch.Tensor。from_numpy共享内存改之前先想清楚。训练循环里不用原地操作。累加 loss 记得.item()。形状不确定就 print别猜。混合精度先测 baseline 再上。einsum 写复杂逻辑可读性优先。7. 我个人在实际操作中的体会写了几年 PyTorch最大的感受是矢量化不是一种技巧而是一种思维方式。刚开始你看到的是一个个样本、一次次循环熟练之后你眼里应该是一块块张量、一个个维度。这个转变一旦完成写出来的代码不仅快而且短、清晰、不容易错。张量创建这块我现在的习惯是数据入口统一用torch.tensor并显式指定 dtype 和 device需要占位就用zeros或empty从 NumPy 过来优先as_tensor。这套组合拳打下来因为创建方式引发的 bug 基本绝迹。最后分享一个小技巧如果你不确定某段代码能不能矢量化先写出循环版本然后问自己“这个循环的每一轮之间有没有依赖”。如果没有依赖几乎一定能改成批量运算如果有依赖比如 RNN 的时间步那就只能保留循环但循环体内的单步运算仍然可以矢量化。把这两层分开看优化思路会清晰很多。