深入理解 PyTorch 中的 torch.stack() 方法：原理、代码与应用

2026-02-08 00:39:30 0条评论 2次阅读 0人点赞

在深度学习和张量运算的日常开发中，我们经常面临着将多个数据源合并、或者是将计算过程中产生的多个中间结果整合在一起的需求。PyTorch 为我们提供了 INLINECODE745b6918 和 INLINECODE4c610923 这两种常用的拼接方法，初学者很容易混淆它们。在这篇文章中，我们将深入探讨 torch.stack() 方法，剖析它的工作原理、它与 INLINECODEc3db98ff 的本质区别，并通过丰富的代码示例展示如何在实际项目中高效地使用它。无论你是正在处理数据增强、多模态输入，还是整合模型输出，理解 INLINECODEb46110bb 都将是你掌握 PyTorch 的关键一步。

什么是 torch.stack()？

简单来说，torch.stack() 的作用是沿着一个新的维度连接（堆叠）一个张量序列。

这里的关键词是“新维度”。与我们熟知的 INLINECODE18822ce2 不同，INLINECODEc9f5cb19 是在现有的维度上进行拼接（例如把两张图片横向拼在一起），而 stack 则是创造一个新的轴，把一系列张量像叠盘子一样叠起来。这就要求所有输入的张量必须具有完全相同的形状。

想象一下，你有两张形状为 [3, 4] 的二维纸。

  • 使用 INLINECODE877a6efd，你可以得到一张更大的纸，形状可能是 INLINECODE2baaee3f（向下拼接）或 [3, 8]（向右拼接）。
  • 使用 INLINECODE01ba079d，你会得到一“摞”纸，形状变成了 INLINECODEe48cf384。这里的 2 就是新增的维度。

语法与参数详解

让我们先来看看它的标准语法：

torch.stack(tensors, dim=0, out=None)

#### 主要参数：

  • tensors (序列): 这是我们想要堆叠的张量序列（比如一个 tuple 或 list）。非常重要的一点是，这些张量必须具有相同的维度和形状。如果你试图把一个 INLINECODEbfefea7d 的张量和一个 INLINECODE2b73c1aa 的张量 stack 在一起，程序会直接报错。
  • dim (整数): 指定我们要在哪一个维度插入这个新轴。它的取值范围是 INLINECODE1b859141 到 INLINECODE2d9c7d77（输入张量的维度数）。这决定了新维度插入的位置。

#### 返回值：

该方法返回一个沿着新维度连接后的新张量，且不会修改原始的张量数据。

代码实战：从一维到多维的深入探索

为了彻底搞懂 stack 的工作机制，让我们通过一系列循序渐进的 Python 3 示例来演示。

#### 示例 1：堆叠一维张量（基础入门）

首先，我们从最简单的一维张量开始。假设我们收集了两组独立的数据点，现在想把它们组合成一个矩阵。

# Python 3 程序演示 torch.stack() 方法针对一维张量的用法
# 导入 torch 库
import torch

# 创建两个一维张量
x = torch.tensor([1., 3., 6., 10.])
y = torch.tensor([2., 7., 9., 13.])

# 打印原始张量以便对比
print("=== 原始张量 ===")
print("Tensor x:", x)
print("Tensor y:", y)

# 情况 1: 沿默认维度 0 堆叠
# dim=0 意味着在最外层（第0维）增加一个维度，把 x 和 y 作为新张量的行
t_dim0 = torch.stack((x, y), dim=0)
print("
=== 沿维度 0 堆叠 (默认) ===")
print(t_dim0)
print("输出形状:", t_dim0.shape) # 将会看到形状是 [2, 4]

# 情况 2: 沿维度 1 堆叠
# dim=1 意味着在列的方向增加一个维度，结果是“转置”堆叠
t_dim1 = torch.stack((x, y), dim=1)
print("
=== 沿维度 1 堆叠 ===")
print(t_dim1)
print("输出形状:", t_dim1.shape) # 将会看到形状是 [4, 2]

输出结果：

=== 原始张量 ===
Tensor x: tensor([ 1.,  3.,  6., 10.])
Tensor y: tensor([ 2.,  7.,  9., 13.])

=== 沿维度 0 堆叠 (默认) ===
tensor([[ 1.,  3.,  6., 10.],
        [ 2.,  7.,  9., 13.]])
输出形状: torch.Size([2, 4])

=== 沿维度 1 堆叠 ===
tensor([[ 1.,  2.],
        [ 3.,  7.],
        [ 6.,  9.],
        [10., 13.]])
输出形状: torch.Size([4, 2])

原理解析：

在这个例子中，输入 INLINECODEacac54c5 和 INLINECODEf49a5896 都是一维向量（长度为 4）。

  • 当我们设置 INLINECODE3273695e 时，PyTorch 创建了一个大小为 2 的新批次维度，并将 INLINECODEa3b2eb1f 放在第 0 行，INLINECODE36db2598 放在第 1 行。结果是一个 INLINECODE3271c7bf 的矩阵。这在深度学习中常用于构建“Batch（批次）”。
  • 当我们设置 INLINECODE92c9b743 时，新维度被插入在列之后。你可以把 INLINECODE0e938789 的第 0 个元素和 INLINECODE58e22282 的第 0 个元素配对，形成新张量的第 0 行。结果是一个 INLINECODE05a6249c 的矩阵。

#### 示例 2：堆叠二维张量（图像处理的基础）

在处理图像数据时，我们通常处理的是二维矩阵（灰度图）或三维张量（RGB 图）。让我们看看如何堆叠两个二维张量。

# Python 3 程序演示 torch.stack() 方法针对二维张量的用法
import torch

# 创建两个 2x3 的二维张量
x = torch.tensor([[1., 3., 6.], [10., 13., 20.]])
y = torch.tensor([[2., 7., 9.], [14., 21., 34.]])

print("=== 输入的二维张量 ===")
print("Tensor x:
", x)
print("Tensor y:
", y)

# 情况 1: 沿维度 0 堆叠
# 这会创建一个“批次”维度，包含两张 2x3 的“图片”
t_dim0 = torch.stack((x, y), dim=0)
print("
=== 沿维度 0 堆叠 ===")
print("Output:
", t_dim0)
print("形状:", t_dim0.shape) # [2, 2, 3]

# 情况 2: 沿维度 1 堆叠
# 新维度插入在第 1 维（行之间）。这会把两个张量的对应行“上下”拼起来。
t_dim1 = torch.stack((x, y), dim=1)
print("
=== 沿维度 1 堆叠 ===")
print("Output:
", t_dim1)
print("形状:", t_dim1.shape) # [2, 2, 3] - 注意虽然形状数字一样，但数据排列不同

# 情况 3: 沿维度 2 堆叠
# 新维度插入在最内层（列之间）。这会把对应的列元素拼起来。
t_dim2 = torch.stack((x, y), dim=2)
print("
=== 沿维度 2 堆叠 ===")
print("Output:
", t_dim2)
print("形状:", t_dim2.shape) # [2, 2, 2] 等等...为什么是 [2, 2, 2]? 让我们解释一下

输出结果：

=== 输入的二维张量 ===
Tensor x:
 tensor([[ 1.,  3.,  6.],
        [10., 13., 20.]])
Tensor y:
 tensor([[ 2.,  7.,  9.],
        [14., 21., 34.]])

=== 沿维度 0 堆叠 ===
Output:
 tensor([[[ 1.,  3.,  6.],
         [10., 13., 20.]],

        [[ 2.,  7.,  9.],
         [14., 21., 34.]]])
形状: torch.Size([2, 2, 3])

=== 沿维度 1 堆叠 ===
Output:
 tensor([[[ 1.,  3.,  6.],
         [ 2.,  7.,  9.]],

        [[10., 13., 20.],
         [14., 21., 34.]]])
形状: torch.Size([2, 2, 3])

=== 沿维度 2 堆叠 ===
Output:
 tensor([[[ 1.,  2.],
         [ 3.,  7.],
         [ 6.,  9.]],

        [[10., 14.],
         [13., 21.],
         [20., 34.]]])
形状: torch.Size([2, 3, 2])

深入解释：

注意观察 dim=2 的输出。

  • 原始张量形状是 (2, 3)（2行3列）。
  • 当 dim=2 时，新维度插入在最后一维（列）之后。
  • 原来 INLINECODE0a471da6 的第 0 行第 0 列是 INLINECODE052cf1a0，INLINECODE24cbaf14 的第 0 行第 0 列是 INLINECODE4e8eaf60。堆叠后，它们变成了新张量第 0 行第 0 列位置的向量 [1., 2.]。
  • 这种操作在处理多通道特征时非常有用，比如将不同卷积层提取的特征图在通道维度上合并。

#### 示例 3：连接超过两个张量（构建批次数据）

在实际训练模型时，我们往往需要一次性处理成百上千张图片。torch.stack() 非常适合将一个 List 中的所有张量合并成一个大的 Batch Tensor。

# Python 3 程序演示堆叠多个张量
import torch

# 假设这是我们从数据集中加载的三个独立样本
x = torch.tensor([1., 3., 6., 10.])
y = torch.tensor([2., 7., 9., 13.])
z = torch.tensor([4., 5., 8., 11.])

print("=== 原始张量 ===")
print("x:", x)
print("y:", y)
print("z:", z)

# 将它们放在一个列表中
tensors_list = [x, y, z]

# 使用 torch.stack 将列表合并为一个 Batch
# dim=0 表示第 0 维现在是 Batch Size
batch_tensor = torch.stack(tensors_list, dim=0)

print("
=== 合并后的 Batch Tensor (dim=0) ===")
print(batch_tensor)
print("形状:", batch_tensor.shape) # [3, 4]，即 3 个样本，每个样本 4 个特征

# 我们也可以尝试 dim=1，这会变成 4 行 3 列的矩阵（转置效果）
batch_dim1 = torch.stack(tensors_list, dim=1)
print("
=== 合并后的 Tensor (dim=1) ===")
print(batch_dim1)
print("形状:", batch_dim1.shape) # [4, 3]

输出结果：

=== 原始张量 ===
x: tensor([ 1.,  3.,  6., 10.])
y: tensor([ 2.,  7.,  9., 13.])
z: tensor([ 4.,  5.,  8., 11.])

=== 合并后的 Batch Tensor (dim=0) ===
tensor([[ 1.,  3.,  6., 10.],
        [ 2.,  7.,  9., 13.],
        [ 4.,  5.,  8., 11.]])
形状: torch.Size([3, 4])

这是构建神经网络输入数据最标准的操作：将 N 个形状为 INLINECODEac3f68f5 的图片，stack 成形状为 INLINECODE0e736979 的张量。

实战应用场景与最佳实践

了解了基础语法后，让我们来看看在实际开发中哪些场景最适合使用 torch.stack。

#### 1. RNN/LSTM 序列处理

在自然语言处理（NLP）中，我们通常将一个句子看作是一个单词序列的列表。假设我们有一个包含 10 个单词的句子，每个单词被转换为一个长度为 100 的词向量（Word Embedding）。我们会有一个包含 10 个 [100] 形状张量的列表。

如果不使用 INLINECODEb33ff854，我们就只能手动处理索引。使用 INLINECODE670d457f 后，我们直接得到一个 [10, 100] 的矩阵，这正是 RNN 期望的输入格式。

# 模拟 NLP 输入
seq_len = 5 # 句子长度
embed_dim = 4 # 词向量维度

# 假设这是 5 个时间步的输入
tensors_sequence = [torch.randn(embed_dim) for _ in range(seq_len)]

# 堆叠成 [Seq_Len, Embed_Dim]
packed_input = torch.stack(tensors_sequence, dim=0)
print("RNN 输入形状:", packed_input.shape) # torch.Size([5, 4])

#### 2. 强化学习中的经验回放

在训练智能体时，我们需要存储“经验”。通常我们会存储 INLINECODEa1566f36。当我们从记忆库中采样出一个 Batch 的经验时，比如 64 条经验，我们会得到 64 个独立的 INLINECODE78a40690 张量。为了将这些数据输入到神经网络进行训练，必须使用 INLINECODE48924e49 将它们转换为 INLINECODE5525cc92 的张量。

常见错误与陷阱

在使用 torch.stack 时，新手（甚至老手）最容易遇到以下报错：

RuntimeError: stack expects each tensor to be equal size

这是因为你试图堆叠形状不一致的张量。

  • 错误示例： torch.stack([tensor([1, 2]), tensor([1, 2, 3])])
  • 原因： 第一个张量长度是 2，第二个是 3。INLINECODEf7d9005e 不像 INLINECODEa006e7a2 那样能处理这种差异。
  • 解决方案： 确保所有输入张量在 INLINECODEbc41ac7c 之前都经过了填充或裁剪，保证形状完全一致。你可以使用 INLINECODE5e65eb9f 来辅助处理不同长度的序列，然后再尝试堆叠。

性能优化与替代方案

虽然 torch.stack 非常方便，但它也有性能开销。因为它需要在内存中重新分配一块连续的内存区域，并将数据从旧的张量复制过去。

  • 优化建议： 如果你在数据加载阶段使用 INLINECODE1799e36b，请确保你的 INLINECODE3e81d5f3 类中的 INLINECODEc9ea2043 方法返回的是正确的形状。PyTorch 的 INLINECODEe9756a74 默认会使用 INLINECODE7ab59db6，其内部实现正是调用了 INLINECODEf5ece465 来自动合并 Batch 数据。因此，大多数情况下你不需要手动写循环来 stack，直接让 DataLoader 去做即可，这样效率最高且支持多进程加载。

总结

在这篇文章中，我们详细探讨了 PyTorch 中 torch.stack() 的用法。让我们回顾一下核心要点：

  • 核心功能： torch.stack() 用于沿新维度连接一系列形状相同的张量。
  • 维度理解： 记住“叠盘子”的比喻。新维度的位置由 dim 参数决定，这决定了数据是如何“堆”在一起的。
  • 形状一致性： 所有输入张量必须形状完全匹配，否则会报错。
  • 实际应用： 它是构建 Batch 数据、处理序列数据和合并模型结果的标准工具。

掌握好 INLINECODEf8572e6b，能让你的数据预处理代码更加简洁、高效。下次当你需要把多个张量合并时，不妨停下来思考一下：我是要增加一个新维度来观察这些数据，还是要在现有维度上把它们拼起来？如果你确定需要前者，INLINECODE0578933f 就是你不二的选择。

希望这篇详细的指南能帮助你更好地理解和使用 PyTorch！继续加油，写出更优雅的深度学习代码吧！

声明：本站所有文章，如无特殊说明或标注，均为本站原创发布。任何个人或组织，在未征得本站同意时，禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益，可联系我们进行处理。如需转载，请注明文章出处豆丁博客和来源网址。https://shluqu.cn/33872.html
点赞
0.00 平均评分 (0% 分数) - 0 票

相关文章美国1G带宽/1T流量高速vps $17.99/年