PyTorch矢量化与张量创建：从显存爆炸到性能优化实战

1. 从一次显存爆炸说起：为什么矢量化值得单独记一笔

去年帮一个朋友排查训练脚本的显存溢出问题，模型本身不大，参数量也就几百万，但一跑起来显存直接飙到 20G 以上。我让他把数据加载和预处理那段代码发过来，扫了一眼就发现问题所在——他在一个for循环里逐条处理样本，每条样本单独做一次张量拼接，循环跑了上万次。这种写法在 CPU 上跑小数据集时看不出毛病，一旦搬到 GPU 上，每次循环都触发一次主机与设备之间的数据搬运，显存碎片化加上中间张量无法及时释放，显存不炸才怪。

这件事让我意识到，矢量化（Vectorization）这个概念虽然听起来基础，但真正在工程实践中踩坑的人不在少数。很多人学 PyTorch 的时候把注意力全放在模型结构、损失函数、优化器上，对张量创建和操作的理解停留在"能用就行"的层面，结果到了实际项目里，性能瓶颈往往就出在这些最基础的地方。

这篇笔记想聊的就是 PyTorch 里矢量化操作和张量创建这两件事。不打算写成 API 手册——那种东西官方文档已经够全了——而是从实际写代码的角度，把那些文档里不会专门讲、但用起来容易出问题的地方捋一遍。适合已经能跑通简单模型、想进一步提升代码质量的读者，也适合刚接触 PyTorch 想少走弯路的新手。核心关键词包括PyTorch、矢量化、张量、torch.tensor、torch.Tensor，这些概念会贯穿全文。

先说结论：矢量化不是"可选优化"，而是 PyTorch 代码的基本功。你写的每一行张量操作，要么在利用底层并行能力，要么在浪费它。区别就在于你有没有意识到这件事。

2. torch.tensor 和 torch.Tensor：一个字母大小写引发的困惑

2.1 两个看起来一样的东西，行为却不同

刚学 PyTorch 的时候，我一度以为torch.tensor和torch.Tensor是同一个东西的两种写法。直到有一次用torch.Tensor(3, 4)想创建一个 3x4 的矩阵，发现里面全是未初始化的垃圾值，才意识到事情没那么简单。

这两个确实容易混淆，但它们的定位完全不同：

torch.tensor是一个工厂函数，它的工作方式是"你给它什么数据，它就根据数据推断出合适的 dtype，然后创建一个包含这些数据的张量"。比如torch.tensor([1, 2, 3])会得到一个 dtype 为torch.int64的张量，torch.tensor([1.0, 2.0])会得到torch.float32。

torch.Tensor则是一个类，它是torch.FloatTensor的别名。当你写torch.Tensor(3, 4)时，Python 会调用这个类的构造函数，参数被解释为形状，创建一个 3x4 的未初始化浮点张量。里面的值是内存里残留的垃圾数据，可能是任意数字。

注意：torch.Tensor创建出来的张量内容是不确定的，千万不要假设它是零或者任何有意义的初始值。如果你需要一个全零的张量，用torch.zeros；需要全一的，用torch.ones。

2.2 为什么官方推荐用 torch.tensor

PyTorch 官方文档里明确建议使用torch.tensor而不是torch.Tensor，原因有几个层面。

第一是类型安全。torch.tensor会根据输入数据推断 dtype，这符合直觉。你给整数列表就得到整数张量，给浮点列表就得到浮点张量。而torch.Tensor永远返回默认浮点类型，如果你需要整数张量，还得额外转换，多此一举。

第二是行为一致性。torch.tensor的行为是可预测的——它总是把参数当作数据。torch.Tensor的行为则取决于参数类型：传数据列表时它当作数据，传整数时它当作形状。这种重载行为容易让人迷惑。

第三是代码可读性。看到torch.tensor([1, 2, 3])，任何人都知道这是在创建一个包含 1、2、3 的张量。看到torch.Tensor(3)，你得想一下这是在创建长度为 3 的未初始化张量，还是在创建包含值 3 的标量。这种歧义在团队协作中是灾难。

我个人的习惯是：永远用torch.tensor，永远不用torch.Tensor。如果需要一个未初始化的张量，明确写torch.empty，这样意图清晰，别人读代码时不会产生误解。

2.3 从数据到张量的几条路径

实际项目中，创建张量的来源多种多样，不同来源对应不同的最佳实践。

从 Python 列表创建是最常见的场景：

import torch # 从整数列表创建，dtype 自动推断为 int64 a = torch.tensor([1, 2, 3]) print(a.dtype) # torch.int64 # 从浮点列表创建，dtype 自动推断为 float32 b = torch.tensor([1.0, 2.0, 3.0]) print(b.dtype) # torch.float32 # 显式指定 dtype，避免推断带来的意外 c = torch.tensor([1, 2, 3], dtype=torch.float32) print(c.dtype) # torch.float32

从 NumPy 数组创建时，torch.tensor会复制数据，而torch.from_numpy会共享内存：

import numpy as np arr = np.array([1.0, 2.0, 3.0]) # 复制数据，修改张量不影响原数组 t1 = torch.tensor(arr) t1[0] = 100.0 print(arr[0]) # 1.0，原数组不变 # 共享内存，修改张量会影响原数组 t2 = torch.from_numpy(arr) t2[0] = 200.0 print(arr[0]) # 200.0，原数组被修改

这个区别在实际项目中很关键。如果你只是想读取 NumPy 数据做推理，用torch.from_numpy可以避免复制开销。但如果你需要对数据进行变换而不想影响原始数据，就必须用torch.tensor做复制。

提示：torch.from_numpy创建的张量和原 NumPy 数组共享同一块内存，这意味着对其中任何一个的修改都会反映到另一个上。在数据预处理流水线中，如果不注意这一点，可能会出现"改了张量结果原始数据也变了"的诡异 bug。

从已有张量创建新张量时，torch.tensor默认会复制数据，而torch.as_tensor会尽量共享内存：

original = torch.tensor([1.0, 2.0, 3.0]) # 复制，新张量和原张量独立 copy = torch.tensor(original) # 共享内存（如果 dtype 和设备一致） shared = torch.as_tensor(original)

torch.as_tensor的行为是"如果输入已经是目标类型的张量，就直接返回它；否则尝试共享内存；实在不行才复制"。这在需要兼容多种输入类型（列表、NumPy 数组、张量）的函数中特别有用。

3. 矢量化到底在优化什么：从循环到并行

3.1 一个循环改写带来的性能差异

先看一段实际代码。假设我们需要计算一个批次中每个样本与一个权重向量的点积，然后加上偏置。

循环版本：

import torch import time batch_size = 10000 feature_dim = 512 data = torch.randn(batch_size, feature_dim) weight = torch.randn(feature_dim) bias = torch.randn(1) # 循环版本 start = time.time() results_loop = torch.empty(batch_size) for i in range(batch_size): results_loop[i] = torch.dot(data[i], weight) + bias loop_time = time.time() - start print(f"循环版本耗时: {loop_time:.4f} 秒")

矢量化版本：

# 矢量化版本 start = time.time() results_vec = data @ weight + bias vec_time = time.time() - start print(f"矢量化版本耗时: {vec_time:.4f} 秒")

在我的机器上，循环版本大约需要 0.8 秒，矢量化版本只需要 0.002 秒左右，差距接近 400 倍。这个差距在 GPU 上会更夸张，因为循环版本无法利用 GPU 的并行计算能力，每次循环都要单独调度。

3.2 底层发生了什么

矢量化之所以快，核心原因在于它把"逐元素操作"转化成了"批量操作"，让底层能够调用优化的 BLAS（基础线性代数子程序）库或者 GPU 的并行计算单元。

用生活化的类比来说：循环版本就像你有一个快递站，每次只处理一个包裹，处理完一个再处理下一个，中间还有大量的调度开销。矢量化版本则是把所有包裹一次性放到传送带上，机器自动分拣，吞吐量完全不是一个量级。

在 CPU 上，矢量化操作会利用 SIMD（单指令多数据）指令集，一条指令同时处理多个数据。在 GPU 上，矢量化操作会启动成千上万个线程并行计算。无论哪种情况，关键都是"让硬件一次处理尽可能多的数据"。

PyTorch 的矢量化操作还有一个额外优势：自动微分友好。当你用矢量化方式写前向传播时，反向传播的计算图也是批量构建的，梯度计算同样可以并行。如果用循环逐样本计算，反向传播时每个样本的梯度单独计算再累加，效率极低。

3.3 哪些操作天然支持矢量化

不是所有操作都能矢量化，但绝大多数常见的数学运算都支持。以下是一些典型例子：

条件筛选的矢量化写法特别值得注意。很多人习惯用if判断，但在张量上应该用布尔索引：

a = torch.randn(1000) # 不推荐：循环判断 result = [] for x in a: if x > 0: result.append(x) result = torch.tensor(result) # 推荐：布尔索引 result = a[a > 0]

布尔索引不仅代码更简洁，而且底层是并行执行的，速度差距在大张量上非常明显。

3.4 广播机制：矢量化的隐形推手

广播（Broadcasting）是矢量化操作中最重要的机制之一，但也是最容易被误解的。简单说，广播允许不同形状的张量进行运算，前提是它们的形状"兼容"。

兼容的规则是：从最后一个维度开始比较，两个维度要么相等，要么其中一个为 1，要么其中一个不存在。满足这些条件就可以广播。

# 形状 (3, 4) 和 (4,) 可以广播 a = torch.randn(3, 4) b = torch.randn(4) c = a + b # b 被广播成 (3, 4) # 形状 (3, 1) 和 (1, 4) 可以广播 d = torch.randn(3, 1) e = torch.randn(1, 4) f = d + e # 结果形状 (3, 4) # 形状 (3, 4) 和 (3,) 不能广播 g = torch.randn(3, 4) h = torch.randn(3) # g + h # 报错

最后那个例子经常让人困惑。(3, 4)和(3,)为什么不能广播？因为从最后一个维度开始比较，4 和 3 不相等，也不满足其中一个为 1 的条件。如果你想让它们相加，需要手动把h扩展成(3, 1)：

h = h.unsqueeze(1) # 形状变成 (3, 1) result = g + h # 现在可以广播成 (3, 4)

提示：广播虽然方便，但不要滥用。过度依赖广播会让代码难以阅读，尤其是当张量维度很多的时候。我个人的原则是：如果广播的逻辑不能一眼看出来，就显式地unsqueeze或expand，让形状变化清晰可见。

4. 张量创建中的 dtype 与 device 陷阱

4.1 dtype 推断的默认规则

torch.tensor的 dtype 推断规则看起来简单，但有几个细节容易踩坑。

Python 整数默认推断为torch.int64，Python 浮点数默认推断为torch.float32。这看起来合理，但问题在于：PyTorch 的默认浮点类型是可以修改的。

import torch # 默认情况下 a = torch.tensor([1.0, 2.0]) print(a.dtype) # torch.float32 # 修改默认浮点类型 torch.set_default_dtype(torch.float64) b = torch.tensor([1.0, 2.0]) print(b.dtype) # torch.float64

这个设置是全局的，会影响之后所有未显式指定 dtype 的张量创建。在团队协作中，如果有人改了默认 dtype 而你不知道，可能会出现"为什么我的模型精度变了"或者"为什么显存占用翻倍了"的问题。

我的建议是：在创建张量时尽量显式指定 dtype，尤其是涉及模型参数和输入数据的时候。多写几个字符，省去后面排查问题的时间。

另一个容易忽略的点是：从 NumPy 数组创建张量时，dtype 会从 NumPy 数组继承。如果 NumPy 数组是float64，创建的张量也是float64，而 PyTorch 模型默认使用float32，直接输入会导致类型不匹配报错。

import numpy as np import torch arr = np.array([1.0, 2.0, 3.0]) # 默认 float64 t = torch.tensor(arr) print(t.dtype) # torch.float64 # 如果模型是 float32，需要转换 t = t.float() # 或者 torch.tensor(arr, dtype=torch.float32)

4.2 device 管理的常见错误

把张量搬到 GPU 上是很常见的操作，但设备管理有几个坑。

第一个坑是混用 CPU 和 GPU 张量。如果两个张量在不同设备上，运算会直接报错：

cpu_tensor = torch.randn(3, 4) gpu_tensor = torch.randn(3, 4).cuda() # cpu_tensor + gpu_tensor # 报错：Expected all tensors to be on the same device

解决方法是统一设备。我习惯在代码开头定义一个device变量，所有张量创建时都指定这个设备：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") a = torch.randn(3, 4, device=device) b = torch.randn(3, 4, device=device) c = a + b # 没问题

第二个坑是模型和数据不在同一设备上。定义了模型之后忘了.to(device)，或者数据加载时忘了指定设备，都会导致报错。更隐蔽的情况是模型的一部分在 GPU 上、一部分在 CPU 上，这种错误往往在训练到一半才暴露出来。

第三个坑是频繁的设备间搬运。有些代码在循环里反复调用.cpu()和.cuda()，每次搬运都有开销。正确的做法是尽量在 GPU 上完成所有计算，只在需要输出结果时才搬回 CPU。

注意：.cuda()和.to("cuda")在功能上等价，但.to()更通用，可以接受设备对象、dtype 等多种参数。推荐统一用.to(device)的写法。

4.3 内存共享与视图操作

PyTorch 里很多操作返回的是视图（view），而不是新的张量。视图和原张量共享内存，修改一个会影响另一个。这个特性在矢量化操作中很常见，也容易引发 bug。

a = torch.randn(4, 4) # view 返回视图，共享内存 b = a.view(16) b[0] = 100.0 print(a[0, 0]) # 100.0，原张量被修改 # reshape 可能返回视图，也可能返回副本 c = a.reshape(2, 8) c[0, 0] = 200.0 print(a[0, 0]) # 200.0，如果 reshape 返回的是视图 # clone 返回副本，不共享内存 d = a.clone() d[0, 0] = 300.0 print(a[0, 0]) # 200.0，原张量不变

view和reshape的区别在于：view要求张量在内存中是连续的，否则会报错；reshape会自动处理非连续的情况，必要时复制数据。如果你确定张量是连续的，用view更高效；如果不确定，用reshape更安全。

另一个常见的视图操作是transpose和permute。它们返回的是视图，但会改变张量的步长（stride），导致张量变成非连续的。这时候如果调用view会报错，需要先.contiguous()：

a = torch.randn(3, 4) b = a.transpose(0, 1) # 形状 (4, 3)，非连续 # b.view(12) # 报错 b = b.contiguous() # 复制成连续内存 b = b.view(12) # 现在可以了

这些细节在写矢量化代码时经常遇到，尤其是涉及维度变换的时候。我的经验是：当你对张量的连续性不确定时，先.contiguous()再操作，虽然多一次复制，但能避免很多莫名其妙的报错。

5. 把矢量化思维落到实际代码里

5.1 数据预处理阶段的矢量化改造

数据预处理是矢量化收益最明显的环节之一。以图像归一化为例，逐张处理的做法是：

# 假设 images 是一个列表，每个元素是 (C, H, W) 的张量 normalized = [] for img in images: img = (img - mean) / std normalized.append(img) normalized = torch.stack(normalized)

矢量化做法是把所有图像堆成一个批次，一次性计算：

# 把所有图像堆成 (N, C, H, W) batch = torch.stack(images) # 一次性归一化，mean 和 std 会自动广播 normalized = (batch - mean) / std

如果mean和std的形状是(C, 1, 1)，广播机制会自动把它们扩展到(N, C, H, W)。这样一次运算就完成了所有图像的归一化，比循环快得多。

文本预处理也是类似。假设你需要把一批文本转成索引序列，循环做法是逐个处理，矢量化做法是先用分词器批量编码，再统一转成张量。

5.2 损失函数计算中的矢量化

自定义损失函数时，矢量化写法尤其重要。以对比学习中的余弦相似度为例，循环版本可能是：

def cosine_similarity_loop(a, b): results = [] for i in range(a.size(0)): dot = torch.dot(a[i], b[i]) norm_a = torch.norm(a[i]) norm_b = torch.norm(b[i]) results.append(dot / (norm_a * norm_b)) return torch.tensor(results)

矢量化版本：

def cosine_similarity_vec(a, b): dot = (a * b).sum(dim=1) norm_a = a.norm(dim=1) norm_b = b.norm(dim=1) return dot / (norm_a * norm_b)

矢量化版本不仅代码更短，而且当a和b是 GPU 张量时，所有计算都在 GPU 上并行完成，速度差距可能是几十倍。

提示：写自定义损失函数时，先问自己"这个操作能不能用张量运算表达"。绝大多数情况下答案都是肯定的。只有在涉及动态控制流（比如根据每个样本的不同条件走不同分支）时，才需要考虑循环或其他方案。

5.3 避免隐式的 Python 循环

有些循环是显式的for，有些则是隐式的。比如列表推导式、map、filter这些，本质上都是循环，只是写法不同。在 PyTorch 代码中，这些都应该尽量避免用在张量上。

# 不推荐：列表推导式 result = torch.tensor([x * 2 for x in tensor]) # 推荐：直接张量运算 result = tensor * 2

另一个常见的隐式循环是torch.stack配合列表推导式：

# 不推荐 results = torch.stack([f(x) for x in tensor_list]) # 如果 f 是逐元素操作，直接矢量化 results = f(tensor_list)

当然，有些场景确实需要循环，比如处理变长序列、实现复杂的控制流等。这时候可以考虑用torch.vmap做自动矢量化，或者用torch.jit.script加速循环。但这些属于进阶话题，基础阶段先把能矢量化的都矢量化了。

5.4 性能验证：怎么确认矢量化真的生效了

写完矢量化代码后，最好验证一下性能提升。最简单的方法是用time.time()计时，但更专业的方式是用torch.cuda.Event做 GPU 计时，或者用torch.profiler做性能分析。

import torch import time # CPU 计时 start = time.time() result = vectorized_operation() end = time.time() print(f"耗时: {end - start:.4f} 秒") # GPU 计时（更准确） start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() result = vectorized_operation() end_event.record() torch.cuda.synchronize() # 等待 GPU 完成 print(f"GPU 耗时: {start_event.elapsed_time(end_event):.4f} 毫秒")

GPU 计时必须用torch.cuda.Event，因为time.time()只能测到 CPU 发起操作的时间，不包括 GPU 实际执行的时间。而且 GPU 操作是异步的，不加torch.cuda.synchronize()的话，计时结果完全不准。

我一般会在优化前后各跑几次，取平均值对比。如果矢量化版本没有明显更快，可能是操作本身太小，矢量化开销占比高，或者瓶颈在别的地方（比如数据加载）。

6. 那些文档不会告诉你的实战细节

6.1 原地操作的双刃剑

PyTorch 支持原地操作（in-place operation），比如a.add_(b)等价于a = a + b但直接修改a的内存。原地操作可以节省内存，但在自动微分场景下可能引发问题。

a = torch.randn(3, requires_grad=True) b = a * 2 b.add_(1) # 原地操作 # b.sum().backward() # 可能报错：a variable needed for gradient computation has been modified

原因是反向传播需要用到前向传播时的中间值，如果这些值被原地修改了，梯度计算就会出错。PyTorch 会检测这种情况并报错，但有时候错误信息不够直观。

我的建议是：在模型的前向传播中避免原地操作，除非你很清楚自己在做什么。在数据预处理和不需要梯度的场景下，原地操作可以放心用。

6.2 张量拼接的效率问题

torch.cat和torch.stack是常用的张量拼接操作，但它们在循环中使用时效率很低。每次拼接都会创建新张量并复制数据，循环 N 次就是 N 次复制。

# 不推荐：循环拼接 result = torch.tensor([]) for i in range(1000): result = torch.cat([result, torch.randn(10)]) # 推荐：先收集到列表，最后一次性拼接 tensors = [] for i in range(1000): tensors.append(torch.randn(10)) result = torch.cat(tensors)

第二种写法只做一次拼接，效率高得多。如果预先知道最终大小，还可以先创建空张量再填充：

result = torch.empty(1000, 10) for i in range(1000): result[i] = torch.randn(10)

但这种方式仍然有循环，最好的做法还是想办法把整个操作矢量化。

6.3 随机数生成的可复现性

做实验时经常需要固定随机种子，保证结果可复现。PyTorch 的随机数生成涉及多个层面：

import torch import numpy as np import random def set_seed(seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False

torch.manual_seed设置 CPU 随机种子，torch.cuda.manual_seed_all设置所有 GPU 的随机种子。如果用了 NumPy 或 Python 的 random 模块，也要一并设置。

cudnn.deterministic = True会让 cuDNN 使用确定性算法，但可能降低性能。cudnn.benchmark = False关闭自动调优，也是为了保证可复现性。这两个设置在生产环境中通常不开，只在需要严格复现的实验中使用。

注意：即使设置了所有种子，在某些情况下结果仍然可能不完全一致，比如使用了非确定性的 GPU 操作、多线程数据加载等。完全可复现是一个需要仔细处理的问题，不是设个种子就万事大吉。

6.4 张量打印的截断问题

调试时经常需要打印张量看看内容，但大张量打印出来会刷屏。PyTorch 默认会截断显示：

a = torch.randn(1000, 1000) print(a) # 只显示部分内容，中间用 ... 省略

如果想看完整内容，可以修改打印选项：

torch.set_printoptions(threshold=float('inf')) # 不截断 torch.set_printoptions(precision=4) # 设置小数位数 torch.set_printoptions(linewidth=120) # 设置每行宽度

但打印大张量本身就很慢，调试时更好的做法是打印统计信息：

print(f"形状: {a.shape}") print(f"均值: {a.mean().item():.4f}") print(f"标准差: {a.std().item():.4f}") print(f"最小值: {a.min().item():.4f}") print(f"最大值: {a.max().item():.4f}")

这样既能了解张量的整体情况，又不会刷屏。

6.5 矢量化不是万能的

最后说一个反直觉的点：矢量化不是所有场景都适用。有些情况下，循环反而更合适。

比如处理变长序列时，每个序列长度不同，强行矢量化需要 padding 到最大长度，可能浪费大量计算。这时候用pack_padded_sequence或者按长度分桶可能更高效。

再比如实现复杂的控制流，比如 beam search、动态规划等，这些算法本身就依赖逐步决策，矢量化反而会让代码难以理解和维护。

我的判断标准是：如果操作是逐元素的、独立的，矢量化几乎总是更好；如果操作涉及样本间的依赖或动态控制流，先考虑算法层面的优化，再考虑矢量化。

7. 从张量创建到矢量化的完整实践路径

把前面聊的内容串起来，形成一个可操作的实践路径。

第一步是统一张量创建方式。在项目里约定只用torch.tensor、torch.zeros、torch.ones、torch.empty这些明确的工厂函数，不用torch.Tensor。创建时尽量显式指定 dtype 和 device，避免依赖默认值。

第二步是识别代码中的循环。把所有对张量的for循环找出来，逐个分析能不能矢量化。判断标准很简单：循环体里的操作是否只依赖当前元素？如果是，就能矢量化。

第三步是用张量运算替换循环。逐元素运算用算术运算符，条件判断用布尔索引，聚合操作用sum、mean、max等，维度变换用view、reshape、permute、unsqueeze。

第四步是验证正确性和性能。矢量化改写后，先用小数据对比循环版本和矢量化版本的输出是否一致，再用大数据测性能提升。如果结果不一致，检查广播规则和维度变换是否正确。

第五步是处理边界情况。空张量、单元素张量、非连续张量、不同 dtype 的张量，这些边界情况在矢量化代码中容易出问题，需要专门测试。

这套流程我在多个项目中用过，基本上能把数据预处理和模型前向传播中的性能瓶颈消除掉。当然，每个项目的情况不同，具体问题具体分析。

我在实际使用中发现，矢量化最大的障碍不是技术难度，而是思维习惯。写了多年 Python 循环的人，看到问题第一反应就是写for，需要刻意练习才能养成"先想能不能矢量化"的条件反射。但一旦养成这个习惯，代码质量和运行效率都会有明显提升。

最后分享一个小技巧：如果你不确定某个操作能不能矢量化，先去 PyTorch 文档里搜一下有没有对应的函数。PyTorch 的函数库非常丰富，绝大多数常见操作都有现成的矢量化实现，自己写循环之前先查一查，往往能省不少事。