张量创建看起来是 PyTorch 里最没技术含量的部分——不就是torch.tensor一把梭吗但我见过太多项目在数据预处理阶段慢得离谱最后定位下来问题就出在张量创建和后续运算没有矢量化上。这篇笔记整理的是我在实际项目里反复验证过的一套思路从张量到底怎么创建最划算到什么时候该用torch.as_tensor而不是torch.tensor再到把 Python 循环改写成矢量化操作之后性能能差出多少。内容偏向实战适合已经能跑通 PyTorch 基础代码、但想让训练和推理真正快起来的人。如果你还在纠结环境搭建或者版本对应这篇不是入门安装教程但里面涉及的 API 行为差异对任何阶段的开发者都有参考价值。1. 张量创建不是只有 torch.tensor 一条路很多人写 PyTorch 代码创建张量就一个写法torch.tensor(data)。能用但未必是最优解。PyTorch 提供了至少七八种创建入口每种背后对应不同的内存行为、类型推断逻辑和拷贝策略。选错了不会报错但会在数据量上来之后悄悄吃掉你的性能。1.1 torch.tensor 与 torch.Tensor 的本质区别这两个名字长得几乎一样但行为完全不同是新手最容易踩的坑之一。torch.tensor是一个工厂函数它会根据输入数据推断 dtype并且总是拷贝数据。你传一个 NumPy 数组进去它会复制一份原数组和新张量互不影响。torch.Tensor是一个类准确说是默认张量类型的别名通常是torch.FloatTensor调用它相当于构造一个默认 float32 类型的张量。它有几个容易让人困惑的行为torch.Tensor(3, 4)创建的是 3x4 的未初始化张量里面是内存里的垃圾值不是零。torch.Tensor([1, 2, 3])才会根据列表内容创建但 dtype 固定为 float32整数也会被转成浮点。torch.Tensor(numpy_array)会共享底层内存在某些版本和条件下行为不如torch.tensor明确。我个人的习惯是永远用torch.tensor不用torch.Tensor。前者语义清晰、类型推断合理、拷贝行为可预期。后者只在需要显式指定默认浮点类型且明确知道自己在干什么的时候才用。import torch import numpy as np # 推荐类型推断 明确拷贝 a torch.tensor([1, 2, 3]) # dtypetorch.int64 b torch.tensor([1.0, 2.0, 3.0]) # dtypetorch.float32 # 不推荐未初始化值是随机的 c torch.Tensor(3, 4) # 垃圾值不是零 # 需要零初始化就用 zeros d torch.zeros(3, 4) # 明确的零提示如果你看到代码里有torch.Tensor(...)且不是从列表构造先怀疑它是不是想写torch.zeros或torch.empty。1.2 从已有数据创建as_tensor、from_numpy 与 tensor 的取舍当你手里已经有一个 NumPy 数组想转成张量有三个选择torch.tensor、torch.as_tensor、torch.from_numpy。它们的核心差异在于是否拷贝内存。函数是否拷贝共享内存适用场景torch.tensor总是拷贝否需要独立副本后续修改互不影响torch.as_tensor尽量不拷贝可能共享只读转换追求性能torch.from_numpy不拷贝是明确知道源数组生命周期可控torch.as_tensor的策略是如果 dtype 和设备已经匹配就直接共享内存如果不匹配比如 NumPy 是 float64PyTorch 默认 float32就拷贝一份并转换类型。这个行为在数据加载管道里非常有用——你可以避免不必要的内存复制。但共享内存是把双刃剑。我踩过一次坑用torch.from_numpy把一个大数组转成张量送进模型然后在外面又修改了原 NumPy 数组结果张量里的数据也跟着变了训练结果莫名其妙。后来排查了半天才定位到内存共享。arr np.array([1.0, 2.0, 3.0], dtypenp.float32) t1 torch.tensor(arr) # 拷贝改 arr 不影响 t1 t2 torch.as_tensor(arr) # 共享dtype 匹配时 t3 torch.from_numpy(arr) # 共享 arr[0] 999.0 print(t1[0]) # tensor(1.) print(t2[0]) # tensor(999.) print(t3[0]) # tensor(999.)实际项目里我的选择逻辑是数据预处理阶段用torch.as_tensor省内存一旦数据要进入需要修改的流程立刻用.clone()断开共享。torch.from_numpy我基本只在确定源数组不会再被碰的时候用。1.3 特殊张量的创建zeros、ones、empty、full 与 arange除了从数据创建PyTorch 还有一批用于生成特定结构张量的函数。这些函数看起来简单但empty和zeros的区别、arange的边界行为都是实际写代码时容易出问题的地方。torch.empty只分配内存不初始化速度最快但里面的值是未定义的。什么时候用当你确定接下来会立刻覆盖所有元素的时候。比如# 先分配再填充避免 zeros 的初始化开销 result torch.empty(1000, 1000) torch.matmul(a, b, outresult)torch.zeros和torch.ones会做初始化有额外开销但安全。torch.full用于填充指定值比先 zeros 再赋值更直接。torch.arange的坑在于它和 Python 的range一样是左闭右开但浮点步长时可能出现精度问题# 看起来应该到 1.0但可能因为浮点误差少一个或多一个 torch.arange(0, 1, 0.1) # 长度可能是 10 也可能是 11 # 需要精确控制时用 linspace torch.linspace(0, 1, 11) # 明确指定元素个数我在做位置编码或者学习率调度的时候凡是涉及浮点序列一律用linspace而不是arange避免因为版本或平台差异导致序列长度不一致。2. 矢量化把 for 循环从代码里赶出去张量创建只是第一步真正决定 PyTorch 代码性能的是你有没有把运算矢量化。所谓矢量化简单说就是用张量整体运算代替逐元素的 Python 循环。PyTorch 底层是 C 和 CUDA 实现的一次张量操作可以在底层并行处理成千上万个元素而 Python 层的 for 循环每次迭代都有解释器开销。2.1 一个真实的性能对比循环 vs 矢量化我拿一个实际场景做过测试给一批样本计算欧氏距离。假设有 10000 个查询向量和 10000 个候选向量每个向量 128 维。先看循环版本import torch import time query torch.randn(10000, 128) candidate torch.randn(10000, 128) # 循环版本双重 for start time.time() result_loop torch.zeros(10000, 10000) for i in range(10000): for j in range(10000): diff query[i] - candidate[j] result_loop[i, j] torch.sqrt((diff ** 2).sum()) print(f循环耗时: {time.time() - start:.2f}s)这个版本在 CPU 上跑基本是分钟级别而且内存占用巨大。再看矢量化版本# 矢量化版本利用广播 start time.time() # (10000, 1, 128) - (1, 10000, 128) - (10000, 10000, 128) diff query.unsqueeze(1) - candidate.unsqueeze(0) result_vec torch.sqrt((diff ** 2).sum(dim-1)) print(f矢量化耗时: {time.time() - start:.2f}s)矢量化版本通常能快几十倍甚至上百倍具体取决于硬件。但这里有个陷阱(10000, 10000, 128)的中间张量占用内存是 10000×10000×128×4 字节大约 51 GB直接爆内存。所以矢量化不是无脑展开还要考虑内存。更稳妥的做法是分块矢量化def batch_distance(query, candidate, chunk1000): results [] for i in range(0, len(query), chunk): q query[i:ichunk].unsqueeze(1) # (chunk, 1, 128) diff q - candidate.unsqueeze(0) # (chunk, 10000, 128) dist torch.sqrt((diff ** 2).sum(dim-1)) results.append(dist) return torch.cat(results, dim0)这样既利用了矢量化又把内存控制在可接受范围内。这个模式在计算注意力矩阵、相似度矩阵的时候非常常用。2.2 广播机制矢量化的核心武器广播broadcasting是矢量化能成立的基础。它的规则是从最后一个维度开始对齐维度大小为 1 的可以扩展到匹配另一个张量维度缺失的相当于大小为 1。理解广播的关键是记住它不实际复制数据只是在计算时逻辑上扩展。这意味着广播本身几乎不占额外内存但一旦参与运算产生新张量内存才会被分配。a torch.randn(3, 1, 5) b torch.randn(1, 4, 5) c a b # 结果形状 (3, 4, 5)这里a在第二维扩展b在第一维扩展最终得到 3×4×5 的结果。如果手动写循环需要三层嵌套用广播一行搞定。广播的常见误用是维度对不齐。比如你想给一个(batch, features)的张量加上一个(features,)的偏置直接加就行因为(features,)会被当作(1, features)广播。但如果你想加的是(batch,)的逐样本权重就需要显式unsqueezex torch.randn(32, 128) bias torch.randn(128) x bias # OK(32,128) (128,) - (32,128) weight torch.randn(32) x weight # 报错或结果不符合预期 x weight.unsqueeze(1) # OK(32,128) (32,1) - (32,128)我见过有人在调试广播问题时反复试形状其实只要记住从右往左对齐这一条大部分问题都能推出来。2.3 用索引和掩码代替条件循环另一类常见的循环是条件判断。比如把张量里所有负数置零ReLU 的手动实现新手会写# 慢逐元素判断 for i in range(x.shape[0]): for j in range(x.shape[1]): if x[i, j] 0: x[i, j] 0矢量化写法x torch.clamp(x, min0) # 方式一 x torch.where(x 0, torch.zeros_like(x), x) # 方式二 x[x 0] 0 # 方式三布尔索引三种方式都比循环快得多。clamp最简洁where最灵活可以同时处理两个分支布尔索引适合原地修改。布尔索引也叫掩码索引在数据筛选里特别有用。比如你有一个(N, 2)的边界框张量想选出面积大于阈值的那些boxes torch.randn(1000, 4) # x1, y1, x2, y2 areas (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) keep areas 0.5 selected boxes[keep] # 直接筛选无需循环这个模式在目标检测的后处理里到处都是用循环写不仅慢代码还长。3. 创建与运算的衔接那些容易忽略的性能细节张量创建和矢量化运算不是两个独立的话题它们在实际代码里是连在一起的。创建方式选得对不对直接影响后续运算能不能矢量化、需不需要额外拷贝。3.1 设备放置创建时就指定别事后搬一个很常见的低效模式是先在 CPU 上创建张量再.to(device)搬到 GPU。如果这个创建发生在训练循环里每次迭代都做一次 CPU→GPU 拷贝开销累积起来很可观。更好的做法是在创建时就指定设备device torch.device(cuda if torch.cuda.is_available() else cpu) # 不推荐先 CPU 再搬 x torch.zeros(1000, 1000) x x.to(device) # 推荐直接创建在目标设备 x torch.zeros(1000, 1000, devicedevice)对于从 NumPy 数组转换的场景torch.as_tensor支持同时指定 device但注意如果指定了不同设备它就必须拷贝arr np.random.randn(1000, 1000).astype(np.float32) t torch.as_tensor(arr, devicedevice) # 会拷贝到 GPU如果这个数组很大且只读可以考虑用 pinned memory 加速传输但那是另一个话题了。3.2 dtype 一致性隐式转换的代价PyTorch 对 dtype 比较严格不同类型运算会报错或者触发类型提升。类型提升本身有开销而且在 GPU 上可能触发额外的 kernel 调用。我遇到过一个案例数据加载器返回 float64 的张量因为 NumPy 默认 float64模型参数是 float32结果每次前向传播都在做隐式类型转换训练速度比预期慢了将近一倍。改成在Dataset里就转成 float32 之后速度恢复正常。# 在数据加载阶段就统一 dtype class MyDataset(torch.utils.data.Dataset): def __getitem__(self, idx): data self.raw[idx] # numpy float64 return torch.as_tensor(data, dtypetorch.float32) # 明确转换创建时指定 dtype 比事后.float()更清晰也更容易在代码审查时发现问题。3.3 原地操作省内存但要看清楚PyTorch 里很多操作有原地版本函数名带下划线比如add_、mul_、zero_。原地操作不分配新内存在内存紧张或者需要累积梯度的时候有用。但原地操作有风险如果这个张量是计算图的一部分原地修改可能破坏反向传播需要的中间值。PyTorch 有时会报错有时不会后者更危险。x torch.randn(3, requires_gradTrue) y x * 2 x.add_(1) # 可能破坏 y 的反向传播 y.sum().backward() # 可能报错或得到错误梯度我的原则是在torch.no_grad()块里或者明确知道张量不参与梯度计算时才用原地操作。训练循环里的参数更新是典型的适用场景因为那本来就在no_grad下做。4. 从创建到实战几个我反复用到的模式前面讲的都是零件这一节把它们组装成实际项目里能直接用的模式。这些模式我在不同的项目里反复写过基本覆盖了张量创建和矢量化的主要应用场景。4.1 数据预处理管道里的张量创建策略数据预处理是张量创建最密集的地方。一个典型管道是读原始数据可能是 NumPy、PIL 图像、CSV→ 转成张量 → 归一化 → 批处理。这里的关键是尽量晚地创建张量尽量早地矢量化。什么意思如果一批数据可以用 NumPy 矢量化处理就先在 NumPy 里做完最后一次性转成张量。因为 NumPy 到 PyTorch 的转换本身有开销频繁的小规模转换不如一次大规模转换。# 不推荐逐样本转张量再拼接 tensors [] for item in raw_data: t torch.tensor(item) tensors.append(t) batch torch.stack(tensors) # 推荐先在 NumPy 层堆叠再一次转换 import numpy as np arr np.stack(raw_data) # NumPy 层完成堆叠 batch torch.as_tensor(arr) # 一次转换如果raw_data里的元素形状一致第二种写法快很多。如果形状不一致那没办法只能逐个处理但可以考虑用torch.nested或者 padding 到统一形状。4.2 注意力矩阵计算中的矢量化技巧注意力机制是矢量化收益最明显的场景之一。标准的缩放点积注意力def scaled_dot_product_attention(Q, K, V, maskNone): # Q: (batch, heads, seq_len, d_k) # K: (batch, heads, seq_len, d_k) # V: (batch, heads, seq_len, d_v) d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) return torch.matmul(attn, V)这里每一步都是矢量化的matmul批量矩阵乘、transpose交换最后两维、masked_fill用掩码填充、softmax沿最后一维归一化。如果用循环实现代码会长十倍速度会慢百倍。掩码的创建本身也值得注意。mask通常是(batch, 1, 1, seq_len)或(batch, 1, seq_len, seq_len)的形状利用广播和scores对齐。创建掩码时用torch.triu、torch.tril这些函数比手动循环快得多# 因果掩码下三角为 1 seq_len 10 causal_mask torch.tril(torch.ones(seq_len, seq_len))4.3 自定义损失函数时的矢量化改写写自定义损失函数是另一个容易写出循环的地方。比如对比损失需要计算每个样本和所有负样本的距离。循环版本# 慢 loss 0 for i in range(batch_size): for j in range(num_negatives): loss torch.exp(torch.dot(anchor[i], negative[j]))矢量化版本# 快 # anchor: (batch, dim), negative: (num_neg, dim) similarity torch.matmul(anchor, negative.T) # (batch, num_neg) loss torch.exp(similarity).sum()如果还涉及正样本通常写成矩阵形式后一次算完。我在实现 InfoNCE 类损失的时候核心就是一次matmul得到所有相似度然后按行做 softmax整个过程没有循环。4.4 张量创建在模型初始化中的注意事项模型参数初始化也涉及张量创建。PyTorch 的nn.Linear等层默认会初始化参数但如果你想自定义初始化需要注意创建方式。import torch.nn as nn class MyLayer(nn.Module): def __init__(self, in_features, out_features): super().__init__() self.weight nn.Parameter(torch.empty(out_features, in_features)) self.bias nn.Parameter(torch.zeros(out_features)) self.reset_parameters() def reset_parameters(self): # 用 Kaiming 初始化原地操作 nn.init.kaiming_uniform_(self.weight, amath.sqrt(5)) # bias 的初始化依赖于 fan_in fan_in self.weight.size(1) bound 1 / math.sqrt(fan_in) nn.init.uniform_(self.bias, -bound, bound)这里用torch.empty创建权重是因为马上会被kaiming_uniform_覆盖用zeros会多一次无用的初始化。bias用zeros是因为uniform_会原地修改但初始值明确一些也无妨。注意nn.init里的函数大多带下划线是原地操作。如果你在创建后需要保留原始值做对比记得先.clone()。5. 调试矢量化代码的实用手段矢量化代码写起来爽但出错的时候不如循环直观。循环可以打断点逐行看矢量化一出错就是形状不匹配或者数值异常。这一节整理几个我常用的调试手段。5.1 形状追踪从右往左对齐形状不匹配是矢量化最常见的错误。我的排查习惯是把涉及的张量形状写下来从右往左对齐看哪一维对不上。a torch.randn(2, 3, 4) b torch.randn(3, 1) # a: (2, 3, 4) # b: (3, 1) - 广播为 (1, 3, 1) - (2, 3, 4) c a b # OK如果对不上PyTorch 会报错并告诉你哪两个形状不兼容。但有时候它不报错而是静默地广播出你意想不到的结果。比如(3,)和(3, 1)相加结果是(3, 3)这通常不是你想要。我的做法是在关键步骤后打印形状print(fQ: {Q.shape}, K: {K.shape}, scores: {scores.shape})在 Jupyter 里开发的时候这个习惯能省很多时间。5.2 数值验证用小规模数据对比循环和矢量化结果写完矢量化代码怎么确认它和循环版本等价用一个小规模输入两个版本都跑一遍对比结果。# 小规模测试 x torch.randn(3, 4) loop_result manual_loop(x) vec_result vectorized(x) assert torch.allclose(loop_result, vec_result, atol1e-6)torch.allclose比直接靠谱因为浮点运算顺序不同可能带来微小误差。atol和rtol根据你的精度要求调整。这个习惯我在实现自定义算子的时候一定会做。先写一个慢但正确的循环版本作为参考再写矢量化版本用随机数据验证两者一致最后才上大规模数据测性能。5.3 性能剖析用 torch.profiler 定位瓶颈如果矢量化之后还是慢可能是瓶颈不在你优化的那段代码。torch.profiler可以告诉你时间花在哪里。from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: for _ in range(10): output model(input) print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))输出会按耗时排序列出每个操作的耗时。如果发现某个aten::copy_或者aten::to占用很高说明有隐式的设备或类型转换那就是优化点。我一般会在训练循环里跑几十步然后看 profiler 结果重点关注三类操作拷贝、类型转换、小 kernel 调用。小 kernel 调用多通常意味着矢量化不彻底有太多细碎的操作。6. 一些版本差异和兼容性提醒PyTorch 版本迭代快张量创建和矢量化的 API 在不同版本间有细微差异。这些差异在跨版本迁移或者看旧代码的时候会造成困惑。6.1 torch.tensor 的 copy 参数早期版本torch.tensor有一个copy参数可以控制是否拷贝。后来这个参数的行为有过调整现在推荐用torch.as_tensor来表达尽量不拷贝的意图。如果你看到旧代码里有torch.tensor(data, copyFalse)可以理解为现在的torch.as_tensor(data)。6.2 默认 dtype 的变化torch.Tensor的默认类型是 float32这个没变。但torch.tensor从整数列表创建时返回 int64从浮点列表创建时返回 float32这个推断逻辑在不同版本间基本一致。需要注意的是如果你依赖默认 dtype在不同平台上可能有差异显式指定 dtype 总是更安全。6.3 设备无关代码的写法写库或者需要跨设备运行的代码时避免硬编码cuda。用torch.device和.to(device)是标准做法。创建张量时如果不知道目标设备可以先创建在 CPU让调用方决定何时搬运。但如果是库内部的临时张量可以用torch.empty_like或者torch.zeros_like来继承已有张量的设备和 dtypedef my_op(x): # 继承 x 的设备和 dtype避免硬编码 result torch.zeros_like(x) # ... return resultzeros_like、ones_like、empty_like这几个函数在写设备无关代码时特别有用它们自动匹配输入张量的属性。6.4 与 ONNX 导出相关的创建注意事项如果模型需要导出成 ONNX张量创建方式会影响导出结果。动态形状、控制流、某些索引操作在 ONNX 里支持有限。我的经验是导出前尽量用静态形状避免在 forward 里根据输入值动态创建不同形状的张量。如果必须动态用torch.onnx.export的dynamic_axes参数声明哪些维度是动态的。另外torch.tensor从 Python 标量创建张量在 ONNX 里通常没问题但从复杂嵌套结构创建可能不被支持。导出前跑一遍torch.onnx.export并检查警告能提前发现大部分问题。7. 把矢量化思维变成肌肉记忆写了这么多核心其实就一句话看到循环就想能不能用张量运算代替。这个思维习惯不是一天养成的我自己的路径是先从数据加载和损失函数这两个地方开始改因为这两处循环最多、收益最明显。改完之后用 profiler 对比看到速度提升就有动力继续改其他地方。张量创建的选择也是同理。一开始可能觉得torch.tensor和torch.as_tensor差不多但当你处理 GB 级别的数据时一次不必要的拷贝就是几秒的差距。这些细节单独看都很小累积起来就是代码质量的差距。最后分享一个我常用的检查清单写完一段涉及张量创建和运算的代码后过一遍创建张量时dtype 和设备是否明确有没有依赖隐式推断从 NumPy 转换时是否真的需要拷贝能不能用as_tensor有没有可以用广播代替的循环有没有可以用布尔索引代替的条件循环中间张量的内存占用是否可控需不需要分块原地操作是否安全张量是否参与梯度计算这个清单不长但每次过一遍能挡掉大部分性能和正确性问题。